From daa73b908ed079a32d959c39b4e89857eb2fa90d Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 26 Mar 2024 09:30:32 +0800 Subject: [PATCH] fix: dtype --- deepmd/pt/utils/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py index 95ffaafd50..ac251648af 100644 --- a/deepmd/pt/utils/utils.py +++ b/deepmd/pt/utils/utils.py @@ -118,7 +118,7 @@ def dict_to_device(sample_dict): sample_dict[key] = [item.to(DEVICE) for item in sample_dict[key]] if isinstance(sample_dict[key], np.float32): sample_dict[key] = ( - torch.ones(1, dtype=torch.float64, device=DEVICE) * sample_dict[key] + torch.ones(1, dtype=torch.float32, device=DEVICE) * sample_dict[key] ) else: if sample_dict[key] is not None: