Commit 45246db3 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH
Browse files

modified the ToTensor transform to properly work with grasyscale images

parent 4b8da253
Pipeline #26536 passed with stage
in 10 minutes and 6 seconds
......@@ -98,6 +98,8 @@ class ToTensor(object):
self.op = transforms.ToTensor()
def __call__(self, sample):
if len(sample['image'].shape) == 2:
sample['image'] = sample['image'][..., numpy.newaxis]
sample['image'] = self.op(sample['image'])
return sample
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment