@Nikronic
Final and working
class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample):
        image, landmarks = sample['image'], np.array(sample['masks'])

        # swap color axis because
        # numpy image: H x W x C
        # torch image: C X H X W
        image = image.transpose(2,0,1)
        print(image.shape)
        image = torch.from_numpy(image).float()
        
        in_transform = transforms.Compose([transforms.Normalize([0.5, 0.5, 0.5],[0.5, 0.5, 0.5])])
        ## discard the transparent, alpha channel (that's the :3)
        image = in_transform(image)[:3,:,:]
        
               
        return {'image': image,'masks': torch.from_numpy(landmarks)}
        #return torch.from_numpy(image).float()
      
transformed_train_dataset = MasksTrainDataset(csv_file='pneumo_input/train/train-rle.csv',
                                           root_dir='pneumo_input/train/images/256/dicom/',
                                           transform=transforms.Compose([
                                               Rescale(224),
                                               ToTensor(),
                                           ]))

for i in range(len(transformed_train_dataset)):
    print(type(transformed_train_dataset))
    print(len(transformed_train_dataset))
    sample = transformed_train_dataset[i]
    

    print(i,sample['image'],sample['masks'])

    if i == 3:
        break
        
train_dataloader = DataLoader(transformed_train_dataset, batch_size=4,
                        shuffle=True, num_workers=4)