diff --git a/src/lightning/pytorch/utilities/model_helpers.py b/src/lightning/pytorch/utilities/model_helpers.py index 44591aa7f4dc1..daf1c400c03df 100644 --- a/src/lightning/pytorch/utilities/model_helpers.py +++ b/src/lightning/pytorch/utilities/model_helpers.py @@ -104,7 +104,7 @@ def _check_mixed_imports(instance: object) -> None: _R_co = TypeVar("_R_co", covariant=True) # return type of the decorated method -class _restricted_classmethod_impl(Generic[_T, _P, _R_co]): +class _restricted_classmethod_impl(Generic[_T, _R_co, _P]): """Drop-in replacement for @classmethod, but raises an exception when the decorated method is called on an instance instead of a class type."""