Skip to content

Commit 81e2831

Browse files
Add parameter type to to_tensor (#8251)
Co-authored-by: Nicolas Hug <[email protected]>
1 parent 6a936e4 commit 81e2831

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

torchvision/transforms/functional.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import numpy as np
99
import torch
1010
from PIL import Image
11+
from PIL.Image import Image as PILImage
1112
from torch import Tensor
1213

1314
try:
@@ -123,7 +124,7 @@ def _is_numpy_image(img: Any) -> bool:
123124
return img.ndim in {2, 3}
124125

125126

126-
def to_tensor(pic) -> Tensor:
127+
def to_tensor(pic: Union[PILImage, np.ndarray]) -> Tensor:
127128
"""Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
128129
This function does not support torchscript.
129130

0 commit comments

Comments
 (0)