bug fixes to_tensor (#686)

* bug fixes to_tensor * tensor checked

bug fixes to_tensor (#686)
* bug fixes to_tensor * tensor checked
6ee98fc6 · surgan12 · Francisco Massa · 0ce81569 · 6ee98fc6
Commit 6ee98fc6 authored Dec 14, 2018 by surgan12 Committed by Francisco Massa Dec 14, 2018
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

torchvision/transforms/functional.py torchvision/transforms/functional.py +3 -0

No files found.
--- a/torchvision/transforms/functional.py
+++ b/torchvision/transforms/functional.py
@@ -51,6 +51,9 @@ def to_tensor(pic):

    if isinstance(pic, np.ndarray):
        # handle numpy array
+        if pic.ndim == 2:
+            pic = pic[:, :, None]
+
        img = torch.from_numpy(pic.transpose((2, 0, 1)))
        # backward compatibility
        if isinstance(img, torch.ByteTensor):