From d43644428c3f4d197b0f76313dd6e8ec4df800fb Mon Sep 17 00:00:00 2001 From: anyangml Date: Tue, 5 Mar 2024 08:18:58 +0000 Subject: [PATCH] fix: UTs --- deepmd/pt/model/atomic_model/linear_atomic_model.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index a0c34bed6e..663cd0338a 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -245,9 +245,7 @@ def remap_atype( ), "The input `atype` cannot be handled by the type_map." type_2_idx = {atp: idx for idx, atp in enumerate(ori_map)} # this maps the atype in the new map to the original map - mapping = torch.tensor( - [type_2_idx[new_map[idx]] for idx in range(len(new_map))] - ).to(atype.device) + mapping = torch.tensor([type_2_idx[new_map[idx]] for idx in range(len(new_map))], device=atype.device) updated_atype = mapping[atype.long()] return updated_atype