Solved by soloupis in post #10 
              
              
                @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…