Fix bug with to_tensor on the input of ndarray of float32 (#415)

test/test_transforms.py

@@ -269,11 +269,16 @@ class Tester(unittest.TestCase):
             output = trans(img)
             assert np.allclose(input_data.numpy(), output.numpy())
 
-            ndarray = np.random.randint(low=0, high=255, size=(height, width, channels))
+            ndarray = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8)
             output = trans(ndarray)
             expected_output = ndarray.transpose((2, 0, 1)) / 255.0
             assert np.allclose(output.numpy(), expected_output)
 
+            ndarray = np.random.rand(height, width, channels).astype(np.float32)
+            output = trans(ndarray)
+            expected_output = ndarray.transpose((2, 0, 1))
+            assert np.allclose(output.numpy(), expected_output)
+
     @unittest.skipIf(accimage is None, 'accimage not available')
     def test_accimage_to_tensor(self):
         trans = transforms.ToTensor()

torchvision/transforms/functional.py

@@ -47,7 +47,10 @@ def to_tensor(pic):
         # handle numpy array
         img = torch.from_numpy(pic.transpose((2, 0, 1)))
         # backward compatibility
+        if isinstance(img, torch.ByteTensor):
             return img.float().div(255)
+        else:
+            return img
 
     if accimage is not None and isinstance(pic, accimage.Image):
         nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32)