From 89a689679feed9ba84898a0da520a699d517d0ec Mon Sep 17 00:00:00 2001 From: Dmitrii Petukhov Date: Thu, 5 Sep 2019 20:01:04 +0200 Subject: [PATCH] Fix make_grid: support any number of channels in tensor --- torchvision/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index f07a3bb4016..d7ddeb538c9 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -74,7 +74,8 @@ def norm_range(t, range): xmaps = min(nrow, nmaps) ymaps = int(math.ceil(float(nmaps) / xmaps)) 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 for y in irange(ymaps): for x in irange(xmaps):