Improve compatibility with the template

env.yaml

@@ -1,4 +1,4 @@
-name: nn-core
+name: nn-template-core
 channels:
   - defaults
   - pytorch

src/nn_core/callbacks.py

@@ -9,6 +9,7 @@
 from pytorch_lightning.loggers import LightningLoggerBase
 
 from nn_core.common import PROJECT_ROOT
+from nn_core.model_logging import NNLogger
 
 pylogger = logging.getLogger(__name__)
 
@@ -28,7 +29,8 @@ def __init__(self, upload: Optional[Dict[str, bool]], logger: Optional[DictConfi
         self.wandb: bool = self.logger_cfg["_target_"].endswith("WandbLogger")
 
     def on_train_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
-        trainer.logger.log_configuration(model=pl_module)
+        if isinstance(trainer.logger, NNLogger):
+            trainer.logger.log_configuration(model=pl_module)
 
         if "wandb_watch" in self.kwargs:
             trainer.logger.wrapped.watch(pl_module, **self.kwargs["wandb_watch"])

src/nn_core/model_logging.py

@@ -21,7 +21,8 @@ def __init__(self, logger: Optional[LightningLoggerBase], storage_dir: str, cfg)
         self.cfg = cfg
 
     def __getattr__(self, item):
-        return getattr(self.wrapped, item)
+        if self.wrapped is not None:
+            return getattr(self.wrapped, item)
 
     @property
     def save_dir(self) -> Optional[str]:
@@ -110,16 +111,16 @@ def log_configuration(
         if isinstance(cfg, DictConfig):
             cfg: Union[Dict[str, Any], argparse.Namespace, DictConfig] = OmegaConf.to_container(cfg, resolve=True)
 
+        # Store the YaML config separately into the wandb dir
+        yaml_conf: str = OmegaConf.to_yaml(cfg=cfg)
+        run_dir: Path = Path(self.run_dir)
+        run_dir.mkdir(exist_ok=True, parents=True)
+        (run_dir / "config.yaml").write_text(yaml_conf)
+
         # save number of model parameters
         cfg[f"{_STATS_KEY}/params_total"] = sum(p.numel() for p in model.parameters())
         cfg[f"{_STATS_KEY}/params_trainable"] = sum(p.numel() for p in model.parameters() if p.requires_grad)
         cfg[f"{_STATS_KEY}/params_not_trainable"] = sum(p.numel() for p in model.parameters() if not p.requires_grad)
 
         # send hparams to all loggers
         self.wrapped.log_hyperparams(cfg)
-
-        # Store the YaML config separately into the wandb dir
-        yaml_conf: str = OmegaConf.to_yaml(cfg=cfg)
-        run_dir: Path = Path(self.run_dir)
-        run_dir.mkdir(exist_ok=True, parents=True)
-        (run_dir / "config.yaml").write_text(yaml_conf)