Only skip _class_path

Lightning-AI · Nov 15, 2024 · ea8608f · ea8608f
1 parent 27d80e2
commit ea8608f
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/src/lightning/pytorch/loggers/utilities.py b/src/lightning/pytorch/loggers/utilities.py
@@ -63,16 +63,14 @@ def _log_hyperparams(trainer: "pl.Trainer") -> None:
     pl_module = trainer.lightning_module
     datamodule_log_hyperparams = trainer.datamodule._log_hyperparams if trainer.datamodule is not None else False
 
-    lightningcli_internal_hparams = {"_class_path", "_instantiator"}
-
     hparams_initial = None
     if pl_module._log_hyperparams and datamodule_log_hyperparams:
         datamodule_hparams = trainer.datamodule.hparams_initial
         lightning_hparams = pl_module.hparams_initial
         inconsistent_keys = []
         for key in lightning_hparams.keys() & datamodule_hparams.keys():
-            if key in lightningcli_internal_hparams:
-                # Skip LightningCLI's internal hparams
+            if key == "_class_path":
+                # Skip LightningCLI's internal hparam
                 continue
             lm_val, dm_val = lightning_hparams[key], datamodule_hparams[key]
             if (
@@ -93,9 +91,9 @@ def _log_hyperparams(trainer: "pl.Trainer") -> None:
     elif datamodule_log_hyperparams:
         hparams_initial = trainer.datamodule.hparams_initial
 
-    # Don't log LightningCLI's internal hparams
+    # Don't log LightningCLI's internal hparam
     if hparams_initial is not None:
-        hparams_initial = {k: v for k, v in hparams_initial.items() if k not in lightningcli_internal_hparams}
+        hparams_initial = {k: v for k, v in hparams_initial.items() if k != "_class_path"}
 
     for logger in trainer.loggers:
         if hparams_initial is not None: