From 033a72a1bed46f342e3af6e1ead54b45b3380627 Mon Sep 17 00:00:00 2001 From: Adam Dangoor Date: Fri, 2 Feb 2024 21:25:41 +0000 Subject: [PATCH] Add parameter type to `to_tensor` --- 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 698942c56af..c7f88d88faf 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.