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 ...@@ -8,6 +8,7 @@ from typing import Any, List, Optional, Tuple, Union
import numpy as np import numpy as np
import torch import torch
from PIL import Image from PIL import Image
from PIL.Image import Image as PILImage
from torch import Tensor from torch import Tensor
try: try:
...@@ -123,7 +124,7 @@ def _is_numpy_image(img: Any) -> bool: ...@@ -123,7 +124,7 @@ def _is_numpy_image(img: Any) -> bool:
return img.ndim in {2, 3} 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. """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
This function does not support torchscript. 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