From 3b721c77cf7b43a330be9817ad1be1400910e621 Mon Sep 17 00:00:00 2001
From: Lysithea <ucaszxy@outlook.com>
Date: Wed, 27 Mar 2024 14:38:42 +0800
Subject: [PATCH] fix get_ntypes

---
 deepmd/pt/model/model/model.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py
index a62050b2d1..9c2154fb2e 100644
--- 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())
+      
\ No newline at end of file