Unverified Commit 9c387588 authored by Philip Meier's avatar Philip Meier Committed by GitHub
Browse files

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

* type annotations for torchvision/utils.py

* add missing annotation for make_grid

* fix annotation for save_image

* mirror PIL annotation for fp
parent e61538cb
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 make_grid(tensor, nrow=8, padding=2,
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:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment