Perf: use fused Adam optimizer

caic99 · Dec 9, 2024 · 1c8ccf6 · 1c8ccf6
1 parent d162d0b
commit 1c8ccf6
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)