fix bug of coord from extended_coord

dptech-corp · Sep 25, 2023 · ec2a61f · ec2a61f
1 parent 9445700
commit ec2a61f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd_pt/model/model/ener.py b/deepmd_pt/model/model/ener.py
@@ -242,7 +242,7 @@ def forward_lower(
                 model_predict['dforce'] = dforce
         # denoise
         else:
-            coord = extended_coord[:nloc]
+            coord = extended_coord[:, :nloc]
             nlist_list = list(torch.split(nlist, self.descriptor.split_sel, -1))
             if not isinstance(self.rcut, list):
                 nnei_mask = nlist != -1