Preserve number of channels when augmenting images

src/ptbench/models/pasa.py

@@ -197,7 +197,8 @@ class PASA(pl.LightningModule):
 
         # Forward pass on the network
         augmented_images = [self.augmentation_transforms(img) for img in images]
-        augmented_images = torch.unsqueeze(torch.cat(augmented_images, 0), 1)
+        # Combine list of augmented images back into a tensor
+        augmented_images = torch.cat(augmented_images, 0).view(images.shape)
         outputs = self(augmented_images)
 
         training_loss = self.criterion(outputs, labels.double())