From a690eaf829bfa2cdc5f54db31140a51c8054fd70 Mon Sep 17 00:00:00 2001 From: Chun Cai Date: Mon, 24 Mar 2025 11:03:51 +0800 Subject: [PATCH] fix: set fused option for Adam optimizer based on device type --- deepmd/pt/train/training.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 3fd0f9a41f..63e0180ace 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -594,7 +594,9 @@ def warm_up_linear(step, warmup_steps): # author: iProzd if self.opt_type == "Adam": self.optimizer = torch.optim.Adam( - self.wrapper.parameters(), lr=self.lr_exp.start_lr, fused=True + self.wrapper.parameters(), + lr=self.lr_exp.start_lr, + fused=False if DEVICE.type == "cpu" else True, ) if optimizer_state_dict is not None and self.restart_training: self.optimizer.load_state_dict(optimizer_state_dict)