Skip to content

GitLab

Commit df557474 authored by Soumith Chintala's avatar Soumith Chintala Committed by GitHub
Browse files

Merge pull request #19 from fmassa/multi_dim 

Add support for images with other than 3 channels
parents 99ace645 a8856b99
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
torchvision/transforms.py
View file @ df557474
...@@ -34,7 +34,7 @@ class ToTensor(object): ...@@ -34,7 +34,7 @@ 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[1], pic.size[0], 3) img = img.view(pic.size[1], pic.size[0], len(pic.mode))
# put it from HWC 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, 1).transpose(0, 2).contiguous() img = img.transpose(0, 1).transpose(0, 2).contiguous()
......
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