Skip to content

GitLab

Commit 6ee98fc6 authored by surgan12's avatar surgan12 Committed by Francisco Massa
Browse files

bug fixes to_tensor (#686) 

* bug fixes to_tensor

* tensor checked
parent 0ce81569
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
torchvision/transforms/functional.py
View file @ 6ee98fc6
...@@ -51,6 +51,9 @@ def to_tensor(pic): ...@@ -51,6 +51,9 @@ def to_tensor(pic):
if isinstance(pic, np.ndarray): if isinstance(pic, np.ndarray):
# handle numpy array # handle numpy array
if pic.ndim == 2:
pic = pic[:, :, None]
img = torch.from_numpy(pic.transpose((2, 0, 1))) img = torch.from_numpy(pic.transpose((2, 0, 1)))
# backward compatibility # backward compatibility
if isinstance(img, torch.ByteTensor): if isinstance(img, torch.ByteTensor):
......
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