From eb68391452624ce5be32951c685d4efdce491b3c Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 20 Mar 2024 03:33:16 -0700 Subject: [PATCH] [fbsync] Add parameter type to `to_tensor` (#8251) Summary: Co-authored-by: Nicolas Hug Reviewed By: vmoens Differential Revision: D55062810 fbshipit-source-id: 142052d28b3d5948876ebf63cc64fd3e2f28e37b --- torchvision/transforms/functional.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py index b5c242c1e64..6578f6c1c67 100644 --- a/torchvision/transforms/functional.py +++ b/torchvision/transforms/functional.py @@ -8,6 +8,7 @@ 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.