fix type hints

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · May 28, 2024 · 88c5219 · 88c5219
1 parent 0bcd410
commit 88c5219
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py
@@ -227,7 +227,7 @@ def get_model_def_script(self) -> str:
         """Get the model definition script."""
         return self.backbone_model.get_model_def_script()
 
-    def get_min_nbor_dist(self) -> float:
+    def get_min_nbor_dist(self) -> Optional[float]:
         """Get the minimum neighbor distance."""
         return self.backbone_model.get_min_nbor_dist()
 

diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py
@@ -135,7 +135,7 @@ def get_model_def_script(self) -> str:
         return self.model.get_model_def_script()
 
     @torch.jit.export
-    def get_min_nbor_dist(self) -> float:
+    def get_min_nbor_dist(self) -> Optional[float]:
         """Get the minimum neighbor distance."""
         return self.model.get_min_nbor_dist()
 

diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py
@@ -292,7 +292,7 @@ def get_model_def_script(self) -> str:
         return self.backbone_model.get_model_def_script()
 
     @torch.jit.export
-    def get_min_nbor_dist(self) -> float:
+    def get_min_nbor_dist(self) -> Optional[float]:
         """Get the minimum neighbor distance."""
         return self.backbone_model.get_min_nbor_dist()