Fix missing best_metric in the temporary PET checkpoints

metatensor · Nov 27, 2024 · 6a7f4c5 · 6a7f4c5
1 parent af0b280
commit 6a7f4c5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/metatrain/experimental/pet/trainer.py b/src/metatrain/experimental/pet/trainer.py
@@ -610,9 +610,11 @@ def train(
                 else:
                     lora_state_dict = None
                 last_model_checkpoint = {
+                    "architecture_name": "experimental.pet",
                     "trainer_state_dict": trainer_state_dict,
                     "model_state_dict": last_model_state_dict,
                     "best_model_state_dict": self.best_model_state_dict,
+                    "best_metric": self.best_metric,
                     "hypers": self.hypers,
                     "epoch": self.epoch,
                     "dataset_info": model.dataset_info,