fix get_ntypes

deepmodeling · Mar 27, 2024 · 3b721c7 · 3b721c7
1 parent 1423750
commit 3b721c7
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py
@@ -45,3 +45,9 @@ def compute_or_load_stat(
     def get_model_def_script(self) -> str:
         """Get the model definition script."""
         return self.model_def_script
+
+    @torch.jit.export
+    def get_ntypes(self):
+        """Returns the number of element types."""
+        return len(self.get_type_map())
+