not extend atype for training

Author: caic99

deepmd/pt/model/descriptor/dpa3.py

@@ -469,16 +469,18 @@ def forward(
         # cast the input to internal precsion
         extended_coord = extended_coord.to(dtype=self.prec)
         nframes, nloc, nnei = nlist.shape
-        nall = extended_coord.view(nframes, -1).shape[1] // 3
-
-        node_ebd_ext = self.type_embedding(extended_atype)
-        node_ebd_inp = node_ebd_ext[:, :nloc, :]
+        # nall = extended_coord.view(nframes, -1).shape[1] // 3
+        if comm_dict is None:
+            atype = extended_atype[:, :nloc]
+        else:
+            atype = extended_atype
+        node_ebd_inp = self.type_embedding(atype)
         # repflows
-        node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows(
+        node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows.forward(
             nlist,
             extended_coord,
             extended_atype,
-            node_ebd_ext,
+            node_ebd_inp,
             mapping,
             comm_dict=comm_dict,
         )

deepmd/pt/model/descriptor/repflows.py

@@ -373,13 +373,12 @@ def forward(
         nlist: torch.Tensor,
         extended_coord: torch.Tensor,
         extended_atype: torch.Tensor,
-        extended_atype_embd: Optional[torch.Tensor] = None,
+        extended_atype_embd: torch.Tensor,
         mapping: Optional[torch.Tensor] = None,
         comm_dict: Optional[dict[str, torch.Tensor]] = None,
     ):
         if comm_dict is None:
             assert mapping is not None
-            assert extended_atype_embd is not None
         nframes, nloc, nnei = nlist.shape
         nall = extended_coord.view(nframes, -1).shape[1] // 3
         atype = extended_atype[:, :nloc]
@@ -402,14 +401,9 @@ def forward(
         # beyond the cutoff sw should be 0.0
         sw = sw.masked_fill(~nlist_mask, 0.0)
 
+        atype_embd = extended_atype_embd
         # [nframes, nloc, tebd_dim]
-        if comm_dict is None:
-            assert isinstance(extended_atype_embd, torch.Tensor)  # for jit
-            atype_embd = extended_atype_embd[:, :nloc, :]
-            assert list(atype_embd.shape) == [nframes, nloc, self.n_dim]
-        else:
-            atype_embd = extended_atype_embd
-        assert isinstance(atype_embd, torch.Tensor)  # for jit
+        assert list(atype_embd.shape) == [nframes, nloc, self.n_dim]
         node_ebd = self.act(atype_embd)
         n_dim = node_ebd.shape[-1]
         # nb x nloc x nnei x 1,  nb x nloc x nnei x 3