update tests

nitrain · May 3, 2024 · 0413ea6 · 0413ea6
1 parent 602d5a1
commit 0413ea6
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/test_integrations.py b/tests/test_integrations.py
@@ -56,6 +56,7 @@ def test_classification(self):
 
         # inference on test data
         predictor = nt.Predictor(model, 
+                                 task='segmentation',
                                  sampler=SliceSampler(batch_size=20, axis=-1))
         #y_pred = predictor.predict(data_test)