Merge pull request #29 from grok-ai/develop

Bump to lighting 2.0.*

setup.cfg

@@ -17,7 +17,7 @@ packages=find:
 install_requires =
     # Add project specific dependencies
     # Stuff easy to break with updates
-    pytorch-lightning==1.7.*
+    lightning==2.0.*
     hydra-core
     wandb
 

src/nn_core/model_logging.py

@@ -10,7 +10,7 @@
 from omegaconf import DictConfig, OmegaConf
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
-from pytorch_lightning.loggers import LightningLoggerBase
+from pytorch_lightning.loggers.logger import Logger
 
 from nn_core.common import PROJECT_ROOT
 
@@ -20,9 +20,9 @@
 _STATS_KEY: str = "stats"
 
 
-class NNLogger(LightningLoggerBase):
+class NNLogger(Logger):
 
-    __doc__ = LightningLoggerBase.__doc__
+    __doc__ = Logger.__doc__
 
     def __init__(self, logging_cfg: DictConfig, cfg: DictConfig, resume_id: Optional[str]):
         super().__init__()
@@ -39,7 +39,7 @@ def __init__(self, logging_cfg: DictConfig, cfg: DictConfig, resume_id: Optional
             self.logging_cfg.logger.mode = "offline"
 
         pylogger.info(f"Instantiating <{self.logging_cfg.logger['_target_'].split('.')[-1]}>")
-        self.wrapped: LightningLoggerBase = hydra.utils.instantiate(
+        self.wrapped: Logger = hydra.utils.instantiate(
             self.logging_cfg.logger,
             version=self.resume_id,
             dir=os.getenv("WANDB_DIR", "."),
@@ -115,7 +115,7 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
 
         This method logs metrics as as soon as it received them. If you want to aggregate
         metrics for one specific `step`, use the
-        :meth:`~pytorch_lightning.loggers.base.LightningLoggerBase.agg_and_log_metrics` method.
+        :meth:`~pytorch_lightning.loggers.base.Logger.agg_and_log_metrics` method.
 
         Args:
             metrics: Dictionary with metric names as keys and measured quantities as values