Do not specify the dimensions in to_tensor pic type hint

pytorch · Jun 27, 2024 · 78d4089 · 78d4089
1 parent 4f56ec2
commit 78d4089
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py
@@ -10,6 +10,7 @@
 from PIL import Image
 from PIL.Image import Image as PILImage
 from torch import Tensor
+from numpy.typing import NDArray
 
 try:
     import accimage
@@ -124,7 +125,7 @@ def _is_numpy_image(img: Any) -> bool:
     return img.ndim in {2, 3}
 
 
-def to_tensor(pic: Union[PILImage, np.ndarray[int, int]]) -> Tensor:
+def to_tensor(pic: Union[PILImage, NDArray[np.uint8]]) -> Tensor:
     """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
     This function does not support torchscript.