check the init args only when the given frames are in __init__ method. (#21227)

* check the init args only when the given frames are in `__init__` method.

* chlog

* add additional testing.

---------

Co-authored-by: jirka <jirka.borovec@seznam.cz>
Co-authored-by: Nicki Skafte Detlefsen <skaftenicki@gmail.com>
Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Co-authored-by: Justus Schock <12886177+justusschock@users.noreply.github.com>
Co-authored-by: Deependu <deependujha21@gmail.com>

Author: 5 people

src/lightning/pytorch/CHANGELOG.md

@@ -61,6 +61,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed `LightningCLI` loading of hyperparameters from `ckpt_path` failing for subclass model mode ([#21246](https://github.com/Lightning-AI/pytorch-lightning/pull/21246))
 
 
+- Fixed check the init args only when the given frames are in `__init__` method ([#21227](https://github.com/Lightning-AI/pytorch-lightning/pull/21227))
+
+
 - Fixed how `ThroughputMonitor` calculated training time ([#21291](https://github.com/Lightning-AI/pytorch-lightning/pull/21291))
 
 

src/lightning/pytorch/utilities/parsing.py

@@ -91,7 +91,7 @@ def get_init_args(frame: types.FrameType) -> dict[str, Any]:  # pragma: no-cover
 
 def _get_init_args(frame: types.FrameType) -> tuple[Optional[Any], dict[str, Any]]:
     _, _, _, local_vars = inspect.getargvalues(frame)
-    if "__class__" not in local_vars:
+    if "__class__" not in local_vars or frame.f_code.co_name != "__init__":
         return None, {}
     cls = local_vars["__class__"]
     init_parameters = inspect.signature(cls.__init__).parameters

tests/tests_pytorch/models/test_hparams.py

@@ -341,6 +341,20 @@ def __init__(obj, *more_args, other_arg=300, **more_kwargs):
         obj.save_hyperparameters()
 
 
+class _MetaType(type):
+    def __call__(cls, *args, **kwargs):
+        instance = super().__call__(*args, **kwargs)  # Create the instance
+        if hasattr(instance, "_after_init"):
+            instance._after_init(**kwargs)  # Call the method if defined
+        return instance
+
+
+class MetaTypeBoringModel(CustomBoringModel, metaclass=_MetaType):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.save_hyperparameters()
+
+
 if _OMEGACONF_AVAILABLE:
 
     class DictConfSubClassBoringModel(SubClassBoringModel):
@@ -365,6 +379,7 @@ class DictConfSubClassBoringModel: ...
         pytest.param(DictConfSubClassBoringModel, marks=RunIf(omegaconf=True)),
         BoringModelWithMixin,
         BoringModelWithMixinAndInit,
+        MetaTypeBoringModel,
     ],
 )
 def test_collect_init_arguments(tmp_path, cls):
@@ -420,6 +435,21 @@ def _raw_checkpoint_path(trainer) -> str:
     return os.path.join(trainer.checkpoint_callback.dirpath, raw_checkpoint_path)
 
 
+def test_collect_init_arguments_in_other_methods():
+    class _ABCModelCreator:
+        def init(self, model, **kwargs) -> LightningModule:
+            self.model = model
+            return self.model
+
+    class ConcreteModelCreator(_ABCModelCreator):
+        def init(self, model=None, **kwargs) -> LightningModule:
+            return super().init(model=model or CustomBoringModel(**kwargs))
+
+    model_creator = ConcreteModelCreator()
+    model = model_creator.init(batch_size=123)
+    assert model.hparams.batch_size == 123
+
+
 @pytest.mark.parametrize("base_class", [HyperparametersMixin, LightningModule, LightningDataModule])
 def test_save_hyperparameters_under_composition(base_class):
     """Test that in a composition where the parent is not a Lightning-like module, the parent's arguments don't get