modified the ToTensor transform to properly work with grasyscale images

bob/learn/pytorch/datasets/utils.py

@@ -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