Fixes #349 (#355)

torchvision/utils.py

@@ -93,9 +93,8 @@ def save_image(tensor, filename, nrow=8, padding=2,
         **kwargs: Other arguments are documented in ``make_grid``.
     """
     from PIL import Image
-    tensor = tensor.cpu()
     grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
                      normalize=normalize, range=range, scale_each=scale_each)
-    ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy()
+    ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
     im = Image.fromarray(ndarr)
     im.save(filename)