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

* fix save_image when height or width == 1

* test_save_image_single_pixel

test/test_utils.py

@@ -43,6 +43,12 @@ class Tester(unittest.TestCase):
             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()

torchvision/utils.py

@@ -67,7 +67,7 @@ def make_grid(tensor, nrow=8, padding=2,
             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)