Add parameter type to to_tensor

pytorch · Feb 5, 2024 · a1cd97c · a1cd97c
1 parent 806dba6
commit a1cd97c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchvision/transforms/functional.py b/torchvision/transforms/functional.py
@@ -123,7 +123,7 @@ def _is_numpy_image(img: Any) -> bool:
     return img.ndim in {2, 3}
 
 
-def to_tensor(pic) -> Tensor:
+def to_tensor(pic: Union[Image, np.ndarray]) -> Tensor:
     """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
     This function does not support torchscript.