Skip to content

Commit

Permalink
Add parameter type to to_tensor (#8251)
Browse files Browse the repository at this point in the history
Co-authored-by: Nicolas Hug <[email protected]>
  • Loading branch information
adamtheturtle and NicolasHug authored Feb 8, 2024
1 parent 6a936e4 commit 81e2831
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion torchvision/transforms/functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit 81e2831

Please sign in to comment.