fix save_image when height or width == 1 (#1059)

* fix save_image when height or width == 1 * test_save_image_single_pixel
pytorch · Jun 26, 2019 · 427633a · 427633a
1 parent 2b6da28
commit 427633a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -43,6 +43,12 @@ def test_save_image(self):
             utils.save_image(t, f.name)
             assert os.path.exists(f.name), 'The image is not present after save'
 
+    def test_save_image_single_pixel(self):
+        with tempfile.NamedTemporaryFile(suffix='.png') as f:
+            t = torch.rand(1, 3, 1, 1)
+            utils.save_image(t, f.name)
+            assert os.path.exists(f.name), 'The pixel image is not present after save'
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/torchvision/utils.py b/torchvision/utils.py
@@ -67,7 +67,7 @@ def norm_range(t, range):
             norm_range(tensor, range)
 
     if tensor.size(0) == 1:
-        return tensor.squeeze()
+        return tensor.squeeze(0)
 
     # make the mini-batch of images into a grid
     nmaps = tensor.size(0)