Tweak probunet input

platipy/imaging/cnn/prob_unet.py

@@ -222,7 +222,7 @@ def __init__(
         self.no_convs_fcomb = no_convs_fcomb
         self.initializers = {"w": "he_normal", "b": "normal"}
         self.z_prior_sample = 0
-        
+
         unet_input_channels = input_channels
         if use_structure_context:
             unet_input_channels += 1

platipy/imaging/cnn/train.py

@@ -579,18 +579,25 @@ def validate(
     def training_step(self, batch, _):
         x, c, y, m, _ = batch
 
+        # Concat input mask if we are using the structure as context
+        if self.use_structure_context:
+            x = torch.cat((x, y), dim=1)
+
+        print(f"y.shape1 {y.shape}")
         # Add background layer for one-hot encoding
         not_y = 1 - y.max(axis=1).values
         not_y = torch.unsqueeze(not_y, dim=1)
         y = torch.cat((not_y, y), dim=1).float()
+        print(f"y.shape2 {y.shape}")
+
+        print(f"x.shape0 {x.shape}")
+
+        print(f"c.shape {c.shape}")
 
         # Concat context map to image if we have one
         if c.numel() > 0:
             x = torch.cat((x, c), dim=1)
-
-        # Concat input mask if we are using the structure as context
-        if self.use_structure_context:
-            x = torch.cat((x, y), dim=1)
+        print(f"x.shape 1 {x.shape}")
 
         # self.prob_unet.forward(x, y, training=True)
         if self.hparams.prob_type == "prob":