Merge branch 'devel' into devel-atomic_weight

deepmodeling · Dec 18, 2024 · a8e204f · a8e204f
2 parents 9d0d7d6 + 104fc36
commit a8e204f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -579,7 +579,7 @@ 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
+                self.wrapper.parameters(), lr=self.lr_exp.start_lr, fused=True
             )
             if optimizer_state_dict is not None and self.restart_training:
                 self.optimizer.load_state_dict(optimizer_state_dict)