Pytorch lightning validation set has different image sizes than training set

Question

When i try to train a cnn, I get different shapes for the same dataloader and i dont know why. This is the output of the shapes I feed into the model:

You can see that my validation shape is [batch size, 1, image height and width]. for some reason, the image size gets changed in the last step and the batch size is 1. The same happens when I use the sanity check from pytorch lightning beforehand, which ive disabled for now. This is how the pytorch lightning data module looks like which gets the dataloader:

class MRIDataModule(pl.LightningDataModule):
    def __init__(self, batch_size, data_paths):
        super().__init__()
        self.batch_size = batch_size
        self.data_paths = data_paths
        self.train_set = None
        self.val_set = None

    def setup(self, stage=None):
        loader = get_data_loader()
        self.train_set = loader(self.data_paths['train_dir'], transform=None, dimension=DIMENSION, nslice=NSLICE)
        self.val_set = loader(self.data_paths['val_dir'], transform=None, dimension=DIMENSION, nslice=NSLICE)

    def train_dataloader(self):
        return DataLoader(self.train_set, batch_size=self.batch_size, num_workers=NUM_WORKERS, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.val_set, batch_size=self.batch_size, num_workers=NUM_WORKERS, shuffle=False)

here is the full code and the print statements are directly from the forward function of my model: https://colab.research.google.com/drive/1yfbCZlwNMqaW1egaTF8HHRD4Ko8iMTxr?usp=sharing