Update CBCT/Model.py

CBCT/Model.py

@@ -37,8 +37,9 @@ class Concat(Reduction):
 
     def __call__(self, tensor: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
         if isinstance(tensor, list):
-            tensor = torch.stack(tensor, dim=0)
-        return tensor.view(tensor.shape[1], -1, *tensor.shape[3:])
+            return torch.stack(tensor, dim=2).squeeze(1)
+        else:
+            return  tensor.view(tensor.shape[0]*tensor.shape[1], -1, *tensor.shape[3:])
         
 class Uncertainty(Transform):