Fix two small bugs

metatensor · Dec 8, 2023 · 98e9436 · 98e9436
1 parent 64bc443
commit 98e9436
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/metatensor/models/soap_bpnn/train.py b/src/metatensor/models/soap_bpnn/train.py
@@ -65,5 +65,5 @@ def train(model, train_dataset, hypers=DEFAULT_TRAINING_HYPERS):
 
     # Save the model:
     save_model(
-        ARCHITECTURE_NAME, model, model_hypers, model.all_species, f"model_{epoch}.pt"
+        ARCHITECTURE_NAME, model, model_hypers, model.all_species, f"model_final.pt"
     )
diff --git a/src/metatensor/models/utils/model_io.py b/src/metatensor/models/utils/model_io.py
@@ -21,6 +21,8 @@ def save_model(
 
         hypers (Dict): The hyperparameters used to train the model.
 
+        all_species (List[int]): The list of all species that the model can handle.
+
         path (str): The path to the file.
     """
     torch.save(