Merge pull request #5 from neptune-ai/correct-checkpoint-namespace

Correct checkpoint namespace
neptune-ai · May 4, 2023 · 4e69f16 · 4e69f16
2 parents bb1b306 + c9b315b
commit 4e69f16
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/neptune_pytorch/impl/__init__.py b/src/neptune_pytorch/impl/__init__.py
@@ -192,7 +192,7 @@ def save_checkpoint(self, checkpoint_name: Optional[str] = None):
             # User is not expected to add extension
             checkpoint_name = checkpoint_name + ".pt"
 
-        safe_upload(self._namespace_handler["model"], checkpoint_name, self.model)
+        safe_upload(self._namespace_handler["model"]["checkpoints"], checkpoint_name, self.model)
 
     def __del__(self):
         # Remove hooks

diff --git a/tests/test_e2e.py b/tests/test_e2e.py
@@ -51,8 +51,8 @@ def test_e2e(model, dataset):
 
     run.wait()
     run.exists(f"{npt_logger.base_namespace}/batch/loss")
-    run.exists(f"{npt_logger.base_namespace}/model/checkpoint_1.pt")
-    run.exists(f"{npt_logger.base_namespace}/model/checkpoint_2.pt")
+    run.exists(f"{npt_logger.base_namespace}/model/checkpoints/checkpoint_1.pt")
+    run.exists(f"{npt_logger.base_namespace}/model/checkpoints/checkpoint_2.pt")
     run.exists(f"{npt_logger.base_namespace}/model/model.pt")
     run.exists(f"{npt_logger.base_namespace}/model/summary")
     run.exists(f"{npt_logger.base_namespace}/model/visualization")