Formatting changes for VolSeg2DPredictor tests

DiamondLightSource · Aug 1, 2022 · 74d366a · 74d366a
1 parent 082696f
commit 74d366a
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tests/test_vol_seg_2d_predictor.py b/tests/test_vol_seg_2d_predictor.py
@@ -1,8 +1,7 @@
 import numpy as np
 import pytest
 import torch
-from volume_segmantics.model.operations.vol_seg_2d_predictor import \
-    VolSeg2dPredictor
+from volume_segmantics.model.operations.vol_seg_2d_predictor import VolSeg2dPredictor
 
 
 @pytest.fixture()
@@ -21,7 +20,9 @@ def test_2d_predictor_init(self, volseg_2d_predictor):
     @pytest.mark.gpu
     @pytest.mark.slow
     def test_predict_single_axis(self, volseg_2d_predictor, rand_int_volume):
-        labels, probs = volseg_2d_predictor._predict_single_axis(rand_int_volume)
+        labels, probs = volseg_2d_predictor._predict_single_axis(
+            rand_int_volume, output_probs=False
+        )
         assert isinstance(labels, np.ndarray)
         assert labels.dtype == np.uint8
         assert probs is None