Update inference.py

inference.py

@@ -95,6 +95,7 @@ def segment_all(wt_model, tc_model, et_model, flair_tensor, t1_tensor, t1ce_tens
     wt_label = torch.sigmoid(wt_output.squeeze(0).squeeze(0)) > 0.5
     tc_label = torch.sigmoid(tc_output.squeeze(0).squeeze(0)) > 0.5
     et_label = torch.sigmoid(et_output.squeeze(0).squeeze(0)) > 0.5
+    print(wt_label.shape, tc_label.shape, et_label.shape)
 
     output = np.zeros((128, 128, 128))
     output[(tc_label == 1) & (wt_label == 0)] = 2