Merge pull request #28 from grok-ai/lucmos-patch-1

Fix (again) NNTemplateCore callback
grok-ai · Mar 20, 2023 · 992a88e · 992a88e
2 parents a85b98e + 99544f7
commit 992a88e
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/src/nn_core/callbacks.py b/src/nn_core/callbacks.py
@@ -51,8 +51,9 @@ def on_train_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> No
     def on_save_checkpoint(
         self, trainer: pl.Trainer, pl_module: pl.LightningModule, checkpoint: Dict[str, Any]
     ) -> None:
-        if trainer.logger is not None:
+        if self._is_nnlogger(trainer):
             trainer.logger.on_save_checkpoint(trainer=trainer, pl_module=pl_module, checkpoint=checkpoint)
-            metadata = getattr(trainer.datamodule, "metadata", None)
-            if metadata is not None:
-                checkpoint[METADATA_KEY] = metadata
+        metadata = getattr(trainer.datamodule, "metadata", None)
+        if metadata is not None:
+            checkpoint[METADATA_KEY] = metadata
+