perf: use torch.embedding for type embedding

Author: caic99

deepmd/pt/model/network/network.py

@@ -286,7 +286,7 @@ def forward(self, atype):
         type_embedding:
 
         """
-        return self.embedding(atype.device)[atype]
+        return torch.embedding(self.embedding(atype.device), atype)
 
     def get_full_embedding(self, device: torch.device):
         """