[fbsync] Add parameter type to to_tensor (#8251)

Summary: Co-authored-by: Nicolas Hug <[email protected]> Reviewed By: vmoens Differential Revision: D55062810 fbshipit-source-id: 142052d28b3d5948876ebf63cc64fd3e2f28e37b
pytorch · Mar 20, 2024 · eb68391 · eb68391
1 parent 11598e1
commit eb68391
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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.