make decorator more generic

Author: pmeier

torchvision/prototype/features/_image.py

@@ -3,7 +3,6 @@
 import warnings
 from typing import Any, Optional, Union, Tuple, cast
 
-import PIL.Image
 import torch
 from torchvision.prototype.utils._internal import StrEnum
 from torchvision.transforms.functional import to_pil_image
@@ -79,7 +78,7 @@ def guess_color_space(data: torch.Tensor) -> ColorSpace:
     def show(self) -> None:
         # TODO: this is useful for developing and debugging but we should remove or at least revisit this before we
         #  promote this out of the prototype state
-        cast(PIL.Image.Image, to_pil_image(make_grid(self.view(-1, *self.shape[-3:])))).show()
+        to_pil_image(make_grid(self.view(-1, *self.shape[-3:]))).show()
 
     def draw_bounding_box(self, bounding_box: BoundingBox, **kwargs: Any) -> Image:
         # TODO: this is useful for developing and debugging but we should remove or at least revisit this before we

torchvision/transforms/functional.py

@@ -3,8 +3,7 @@
 import numbers
 import warnings
 from enum import Enum
-from typing import Callable
-from typing import List, Tuple, Any, Optional
+from typing import List, Tuple, Any, Optional, TypeVar, Callable, cast
 
 import numpy as np
 import torch
@@ -20,14 +19,16 @@
 from . import functional_pil as F_pil
 from . import functional_tensor as F_t
 
+F = TypeVar("F", bound=Callable[..., Any])
 
-def log_api_usage_once(fn: Callable[..., Tensor]) -> Callable[..., Tensor]:
+
+def log_api_usage_once(fn: F) -> F:
     @functools.wraps(fn)
-    def wrapper(*args: Any, **kwargs: Any) -> Tensor:
+    def wrapper(*args, **kwargs):
         _log_api_usage_once(fn)
         return fn(*args, **kwargs)
 
-    return wrapper
+    return cast(F, wrapper)
 
 
 class InterpolationMode(Enum):