refactored _restricted_classmethod_impl

Author: GdoongMathew

src/lightning/pytorch/utilities/model_helpers.py

@@ -104,29 +104,23 @@ def _check_mixed_imports(instance: object) -> None:
 _R_co = TypeVar("_R_co", covariant=True)  # return type of the decorated method
 
 
-class _restricted_classmethod_impl(classmethod):
+def _restricted_classmethod_impl(method: Callable[Concatenate[type[_T], _P], _R_co]) -> classmethod:
     """Drop-in replacement for @classmethod, but raises an exception when the decorated method is called on an instance
     instead of a class type."""
 
-    def __init__(self, method: Callable[Concatenate[type[_T], _P], _R_co]) -> None:
-        super().__init__(method)
-        self.method = method
-
-    def __get__(self, instance: _T, cls: Optional[type[_T]] = None) -> Callable[_P, _R_co]:
-        # The wrapper ensures that the method can be inspected, but not called on an instance
-        @functools.wraps(self.method)
-        def wrapper(*args: Any, **kwargs: Any) -> _R_co:
-            # Workaround for https://github.com/pytorch/pytorch/issues/67146
-            is_scripting = any(os.path.join("torch", "jit") in frameinfo.filename for frameinfo in inspect.stack())
-            cls_type = cls if cls is not None else type(instance)
-            if instance is not None and not is_scripting:
-                raise TypeError(
-                    f"The classmethod `{cls_type.__name__}.{self.method.__name__}` cannot be called on an instance."
-                    " Please call it on the class type and make sure the return value is used."
-                )
-            return self.method(cls_type, *args, **kwargs)
-
-        return wrapper
+    # The wrapper ensures that the method can be inspected, but not called on an instance
+    @functools.wraps(method)
+    def wrapper(cls: type[_T], *args: _P.args, **kwargs: _P.kwargs) -> _R_co:
+        # Workaround for https://github.com/pytorch/pytorch/issues/67146
+        is_scripting = any(os.path.join("torch", "jit") in frameinfo.filename for frameinfo in inspect.stack())
+        if inspect.isclass(cls) and not is_scripting:
+            raise TypeError(
+                f"The classmethod `{cls.__name__}.{method.__name__}` cannot be called on an instance."
+                " Please call it on the class type and make sure the return value is used."
+            )
+        return method(cls, *args, **kwargs)
+
+    return classmethod(wrapper)
 
 
 if TYPE_CHECKING: