Implemented native PET training (#349)

metatensor · Oct 8, 2024 · 917a7a7 · 917a7a7
1 parent 161fd56
commit 917a7a7
Show file tree

Hide file tree

Showing 6 changed files with 629 additions and 80 deletions.
diff --git a/src/metatrain/experimental/pet/default-hypers.yaml b/src/metatrain/experimental/pet/default-hypers.yaml
@@ -54,4 +54,5 @@ training:
   DO_GRADIENT_CLIPPING: False
   GRADIENT_CLIPPING_MAX_NORM: null # must be overwritten if DO_GRADIENT_CLIPPING is True
   USE_SHIFT_AGNOSTIC_LOSS: False # only used when fitting general target. Primary use case: EDOS
-  ENERGIES_LOSS: per_structure # per_structure or per_atom
+  ENERGIES_LOSS: per_structure # per_structure or per_atom
+  CHECKPOINT_INTERVAL: 100
diff --git a/src/metatrain/experimental/pet/schema-hypers.json b/src/metatrain/experimental/pet/schema-hypers.json
@@ -219,6 +219,9 @@
         },
         "BALANCED_DATA_LOADER": {
           "type": "boolean"
+        },
+        "CHECKPOINT_INTERVAL": {
+          "type": "integer"
         }
       },
       "additionalProperties": false,