Skip to content

GitLab

Commit cabca398 authored by Dmitrii Petukhov's avatar Dmitrii Petukhov Committed by Francisco Massa
Browse files

Fix make_grid: support any number of channels in tensor (#1300)

parent 7f7e7663
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
torchvision/utils.py
View file @ cabca398
...@@ -74,7 +74,8 @@ def make_grid(tensor, nrow=8, padding=2, ...@@ -74,7 +74,8 @@ def make_grid(tensor, nrow=8, padding=2,
xmaps = min(nrow, nmaps) xmaps = min(nrow, nmaps)
ymaps = int(math.ceil(float(nmaps) / xmaps)) ymaps = int(math.ceil(float(nmaps) / xmaps))
height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding) height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding)
grid = tensor.new_full((3, height * ymaps + padding, width * xmaps + padding), pad_value) num_channels = tensor.size(1)
grid = tensor.new_full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
k = 0 k = 0
for y in irange(ymaps): for y in irange(ymaps):
for x in irange(xmaps): for x in irange(xmaps):
......
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