fix code

HydrogenSulfate · Oct 9, 2024 · f7f64b1 · f7f64b1
1 parent a83fb63
commit f7f64b1
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 8 deletions.
diff --git a/deepmd/pd/loss/dos.py b/deepmd/pd/loss/dos.py
@@ -140,7 +140,7 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
                 [-1, self.numb_dos]
             )
             if "mask" in model_pred:
-                diff = diff[model_pred["mask"].reshape([-1]).bool()]
+                diff = diff[model_pred["mask"].reshape([-1]).astype("bool")]
             l2_local_loss_dos = paddle.mean(paddle.square(diff))
             if not self.inference:
                 more_loss["l2_local_dos_loss"] = self.display_if_exist(
@@ -164,7 +164,7 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
                 [-1, self.numb_dos]
             )
             if "mask" in model_pred:
-                diff = diff[model_pred["mask"].reshape([-1]).bool()]
+                diff = diff[model_pred["mask"].reshape([-1]).astype("bool")]
             l2_local_loss_cdf = paddle.mean(paddle.square(diff))
             if not self.inference:
                 more_loss["l2_local_cdf_loss"] = self.display_if_exist(

diff --git a/deepmd/pd/loss/tensor.py b/deepmd/pd/loss/tensor.py
@@ -107,7 +107,7 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
                 [-1, self.tensor_size]
             )
             if "mask" in model_pred:
-                diff = diff[model_pred["mask"].reshape([-1]).bool()]
+                diff = diff[model_pred["mask"].reshape([-1]).astype("bool")]
             l2_local_loss = paddle.mean(paddle.square(diff))
             if not self.inference:
                 more_loss[f"l2_local_{self.tensor_name}_loss"] = self.display_if_exist(
@@ -133,9 +133,12 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
             if "mask" in model_pred:
                 atom_num = model_pred["mask"].sum(-1, keepdim=True)
                 l2_global_loss = paddle.mean(
-                    paddle.sum(paddle.square(diff) * atom_num, axis=0) / atom_num.sum()
+                    paddle.sum(
+                        paddle.square(diff) * atom_num.astype(diff.dtype), axis=0
+                    )
+                    / atom_num.sum()
                 )
-                atom_num = paddle.mean(atom_num.float())
+                atom_num = paddle.mean(atom_num.astype(diff.dtype))
             else:
                 atom_num = natoms
                 l2_global_loss = paddle.mean(paddle.square(diff))

diff --git a/deepmd/pd/utils/nlist.py b/deepmd/pd/utils/nlist.py
@@ -100,8 +100,7 @@ def build_neighbor_list(
     nall = coord.shape[1] // 3
     # fill virtual atoms with large coords so they are not neighbors of any
     # real atom.
-    # if coord.numel().item() > 0:
-    if True > 0:
+    if coord.numel() > 0:
         xmax = paddle.max(coord) + 2.0 * rcut
     else:
         xmax = paddle.zeros([], dtype=coord.dtype).to(device=coord.place) + 2.0 * rcut

diff --git a/deepmd/pd/utils/preprocess.py b/deepmd/pd/utils/preprocess.py
@@ -183,7 +183,7 @@ def build_neighbor_list(
     - atype: shape is [nall]
     """
     nall = coord.numel() // 3
-    coord = coord.float()
+    coord = coord.astype(paddle.get_default_dtype())
     nlist = [[] for _ in range(nloc)]
     coord_l = coord.reshape([-1, 1, 3])[:nloc]
     coord_r = coord.reshape([1, -1, 3])