Update Lightning to 1.7.*

Author: lucmos

setup.cfg

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

src/nn_core/serialization.py

@@ -12,6 +12,7 @@
 
 import pytorch_lightning as pl
 import torch
+from pytorch_lightning.core.saving import _load_state
 from pytorch_lightning.plugins import TorchCheckpointIO
 
 METADATA_KEY: str = "metadata"
@@ -94,7 +95,7 @@ def remove_checkpoint(self, path) -> None:
 
 
 def compress_checkpoint(src_dir: Path, dst_file: Path, delete_dir: bool = True):
-
+    dst_file.parent.mkdir(parents=True, exist_ok=True)
     with zipfile.ZipFile(_normalize_path(dst_file), "w") as zip_file:
         for folder, subfolders, files in os.walk(src_dir):
             folder: Path = Path(folder)
@@ -161,7 +162,7 @@ def load_model(
         if substitute_values is not None:
             checkpoint = _substistute(checkpoint, substitute_values=substitute_values, substitute_keys=substitute_keys)
 
-        return module_class._load_model_state(checkpoint=checkpoint, metadata=checkpoint.get("metadata", None))
+        return _load_state(cls=module_class, checkpoint=checkpoint, metadata=checkpoint.get("metadata", None))
     else:
         pylogger.warning(f"Loading a legacy checkpoint (from vanilla PyTorch Lightning): '{checkpoint_path}'")
         module_class.load_from_checkpoint(checkpoint_path=str(checkpoint_path), map_location=map_location)