Added packages to pth checkpoint (#1581)

* Added packages to pth checkpoint * Fixed e2e test, added sorted list --------- Co-authored-by: Louis-Dupont <35190946+Louis-Dupont@users.noreply.github.com>
Deci-AI · Oct 29, 2023 · 827c4f8 · 827c4f8
1 parent c0594a3
commit 827c4f8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/super_gradients/training/sg_trainer/sg_trainer.py b/src/super_gradients/training/sg_trainer/sg_trainer.py
@@ -660,7 +660,7 @@ def _save_checkpoint(
         metric = validation_results_dict[self.metric_to_watch]
 
         # BUILD THE state_dict
-        state = {"net": unwrap_model(self.net).state_dict(), "acc": metric, "epoch": epoch}
+        state = {"net": unwrap_model(self.net).state_dict(), "acc": metric, "epoch": epoch, "packages": get_installed_packages()}
 
         if optimizer is not None:
             state["optimizer_state_dict"] = optimizer.state_dict()

diff --git a/tests/end_to_end_tests/trainer_test.py b/tests/end_to_end_tests/trainer_test.py
@@ -79,7 +79,7 @@ def test_checkpoint_content(self):
         ckpt_paths = [os.path.join(trainer.checkpoints_dir_path, suf) for suf in ckpt_filename]
         for ckpt_path in ckpt_paths:
             ckpt = torch.load(ckpt_path)
-            self.assertListEqual(["net", "acc", "epoch", "optimizer_state_dict", "scaler_state_dict"], list(ckpt.keys()))
+            self.assertListEqual(sorted(["net", "acc", "epoch", "optimizer_state_dict", "scaler_state_dict", "packages"]), sorted(list(ckpt.keys())))
         trainer._save_checkpoint()
         weights_only = torch.load(os.path.join(trainer.checkpoints_dir_path, "ckpt_latest_weights_only.pth"))
         self.assertListEqual(["net"], list(weights_only.keys()))