Type annotations for torchvision/utils.py (#2034) (#2428)

Summary:
* type annotations for torchvision/utils.py

* add missing annotation for make_grid

* fix annotation for save_image

* mirror PIL annotation for fp

Pull Request resolved: #2428

Reviewed By: zhangguanheng66

Differential Revision: D22437423

Pulled By: fmassa

fbshipit-source-id: 77a46deb22e4fbf4de02f9f7d0e418c656d40a65

Author: pmeier

torchvision/utils.py

@@ -1,10 +1,15 @@
+from typing import Union, Optional, Sequence, Tuple, Text, BinaryIO
+import io
+import pathlib
 import torch
 import math
 irange = range
 
 
-def make_grid(tensor, nrow=8, padding=2,
-              normalize=False, range=None, scale_each=False, pad_value=0):
+def make_grid(tensor: Union[torch.Tensor, Sequence[torch.Tensor]], nrow: int = 8,
+              padding: int = 2, normalize: bool = False,
+              range: Optional[Tuple[int, int]] = None, scale_each: bool = False,
+              pad_value: int = 0) -> torch.Tensor:
     """Make a grid of images.
 
     Args:
@@ -88,8 +93,9 @@ def norm_range(t, range):
     return grid
 
 
-def save_image(tensor, fp, nrow=8, padding=2,
-               normalize=False, range=None, scale_each=False, pad_value=0, format=None):
+def save_image(tensor: Union[torch.Tensor, Sequence[torch.Tensor]], fp: Union[Text, pathlib.Path, BinaryIO],
+               nrow: int = 8, padding: int = 2, normalize: bool = False, range: Optional[Tuple[int, int]] = None,
+               scale_each: bool = False, pad_value: int = 0, format: Optional[str] = None) -> None:
     """Save a given Tensor into an image file.
 
     Args: