fix(pt): store min_nbor_dist in the state dict

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/pt/entrypoints/main.py

@@ -346,10 +346,14 @@ def train(
     # save min_nbor_dist
     if min_nbor_dist is not None:
         if not multi_task:
-            trainer.model.min_nbor_dist = min_nbor_dist
+            trainer.model.min_nbor_dist = torch.tensor(
+                min_nbor_dist, dtype=torch.float64, device=DEVICE
+            )
         else:
             for model_item in min_nbor_dist:
-                trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item]
+                trainer.model[model_item].min_nbor_dist = torch.tensor(
+                    min_nbor_dist[model_item], dtype=torch.float64, device=DEVICE
+                )
     trainer.run()
 
 

deepmd/pt/model/model/model.py

@@ -18,7 +18,7 @@ def __init__(self, *args, **kwargs):
         """Construct a basic model for different tasks."""
         torch.nn.Module.__init__(self)
         self.model_def_script = ""
-        self.min_nbor_dist = None
+        self.register_buffer("min_nbor_dist", None)
 
     def compute_or_load_stat(
         self,
@@ -50,7 +50,9 @@ def get_model_def_script(self) -> str:
     @torch.jit.export
     def get_min_nbor_dist(self) -> Optional[float]:
         """Get the minimum distance between two atoms."""
-        return self.min_nbor_dist
+        if self.min_nbor_dist is None:
+            return None
+        return self.min_nbor_dist.item()
 
     @torch.jit.export
     def get_ntypes(self):