Skip to content

GitLab

Commit 0a04aaa4 authored by Soumith Chintala's avatar Soumith Chintala
Browse files

fixing ToTensor bugs

parent 386a91b0
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
torchvision/transforms.py
View file @ 0a04aaa4
...@@ -33,10 +33,10 @@ class ToTensor(object): ...@@ -33,10 +33,10 @@ class ToTensor(object):
else: else:
# handle PIL Image # handle PIL Image
img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes())) img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
img = img.view(pic.size[0], pic.size[1], 3) img = img.view(pic.size[1], pic.size[0], 3)
# put it from WHC to CHW format # put it from HWC to CHW format
# yikes, this transpose takes 80% of the loading time/CPU # yikes, this transpose takes 80% of the loading time/CPU
img = img.transpose(0, 2).contiguous() img = img.transpose(0, 1).transpose(0, 2).contiguous()
return img.float().div(255) return img.float().div(255)
class ToPILImage(object): class ToPILImage(object):
......
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