Add missing unsqueeze

pyplati · Dec 6, 2023 · 0ddb078 · 0ddb078
1 parent c8ec57e
commit 0ddb078
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/platipy/imaging/cnn/train.py b/platipy/imaging/cnn/train.py
@@ -605,8 +605,10 @@ def validation_step(self, batch, _):
                 )
                 np.save(mask_file, y[s].cpu().numpy())
 
-            # Image will be same for all in batch
+            # Image (and context map) will be same for all in batch
             x = x[0].unsqueeze(0)
+            if c is not None:
+                c = c[0].unsqueeze(0)
             if self.hparams.ndims == 2:
                 vis = ImageVisualiser(sitk.GetImageFromArray(x.to("cpu")[0]), axis="z")
             else: