Skip to content

Commit e7ed325

Browse files
committed
fix mypy error.
1 parent f21d406 commit e7ed325

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

src/lightning/pytorch/utilities/model_helpers.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -112,18 +112,19 @@ def __init__(self, method: Callable[Concatenate[type[_T], _P], _R_co]) -> None:
112112
super().__init__(method)
113113
self.method = method
114114

115-
def __get__(self, instance: Optional[_T], cls: type[_T]) -> Callable[_P, _R_co]:
115+
def __get__(self, instance: _T, cls: Optional[type[_T]] = None) -> Callable[_P, _R_co]:
116116
# The wrapper ensures that the method can be inspected, but not called on an instance
117117
@functools.wraps(self.method)
118118
def wrapper(*args: Any, **kwargs: Any) -> _R_co:
119119
# Workaround for https://github.com/pytorch/pytorch/issues/67146
120120
is_scripting = any(os.path.join("torch", "jit") in frameinfo.filename for frameinfo in inspect.stack())
121+
cls_type = cls if cls is not None else type(instance)
121122
if instance is not None and not is_scripting:
122123
raise TypeError(
123-
f"The classmethod `{cls.__name__}.{self.method.__name__}` cannot be called on an instance."
124+
f"The classmethod `{cls_type.__name__}.{self.method.__name__}` cannot be called on an instance."
124125
" Please call it on the class type and make sure the return value is used."
125126
)
126-
return self.method(cls, *args, **kwargs)
127+
return self.method(cls_type, *args, **kwargs)
127128

128129
return wrapper
129130

0 commit comments

Comments
 (0)