use log api from #5424

Author: kazhang

torchvision/utils.py

@@ -3,7 +3,7 @@
 import warnings
 from functools import wraps
 from types import FunctionType
-from typing import Any, BinaryIO, List, Optional, Tuple, Union
+from typing import Any, BinaryIO, List, Optional, Tuple, Union, TypeVar, Callable, cast
 
 import numpy as np
 import torch
@@ -549,16 +549,14 @@ def _log_api_usage_once(obj: Any) -> None:
         name = obj.__name__
     torch._C._log_api_usage_once(f"{obj.__module__}.{name}")
 
-def log_api_usage_once_dec(f):
-    @wraps(f)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def log_api_usage_once_dec(fn: F) -> F:
+    @wraps(fn)
     def wrapper(*args, **kwargs):
-        event = f.__module__
-        if f.__name__.endswith("__init__"):
-            # inside module instantiation
-            event += args[0].__class__.__name__
-        else:
-            event += f.__name__
-        torch._C._log_api_usage_once(event)
-        return f(*args, **kwargs)
+        _log_api_usage_once(fn)
+        return fn(*args, **kwargs)
 
-    return wrapper
+    return cast(F, wrapper)