Unverified Commit 81e2831a authored by Adam Dangoor's avatar Adam Dangoor Committed by GitHub
Browse files

Add parameter type to `to_tensor` (#8251)


Co-authored-by: default avatarNicolas Hug <nh.nicolas.hug@gmail.com>
parent 6a936e48
......@@ -8,6 +8,7 @@ from typing import Any, List, Optional, Tuple, Union
import numpy as np
import torch
from PIL import Image
from PIL.Image import Image as PILImage
from torch import Tensor
try:
......@@ -123,7 +124,7 @@ def _is_numpy_image(img: Any) -> bool:
return img.ndim in {2, 3}
def to_tensor(pic) -> Tensor:
def to_tensor(pic: Union[PILImage, np.ndarray]) -> Tensor:
"""Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
This function does not support torchscript.
......
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