fix tuple error (#77)

teamtomo · Jul 12, 2024 · 14157e9 · 14157e9
1 parent 77f8b03
commit 14157e9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/membrain_seg/segmentation/segment.py b/src/membrain_seg/segmentation/segment.py
@@ -159,7 +159,9 @@ def segment(
             with torch.cuda.amp.autocast():
                 mirrored_input = get_mirrored_img(new_data.clone(), m).to(device)
                 mirrored_pred = inferer(mirrored_input, pl_model)
-                if not isinstance(mirrored_pred, list):
+                if not (
+                    isinstance(mirrored_pred, list) or isinstance(mirrored_pred, tuple)
+                ):
                     mirrored_pred = [mirrored_pred]
                 correct_pred = get_mirrored_img(mirrored_pred[0], m)
                 predictions += correct_pred.detach().cpu()