Update training.py

iProzd · iProzd · commit 6e20c181d1ee · 2025-05-21T23:04:31.000+08:00
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -622,6 +622,7 @@ def warm_up_linear(step, warmup_steps):
                     self.wrapper.parameters(),
                     lr=self.lr_exp.start_lr,
                     weight_decay=float(self.opt_param["weight_decay"]),
+                    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)

Original file line number	Diff line number	Diff line change
`@@ -622,6 +622,7 @@ def warm_up_linear(step, warmup_steps):`
`622`	`622`	`self.wrapper.parameters(),`
`623`	`623`	`lr=self.lr_exp.start_lr,`
`624`	`624`	`weight_decay=float(self.opt_param["weight_decay"]),`
	`625`	`+ fused=False if DEVICE.type == "cpu" else True,`
`625`	`626`	`)`
`626`	`627`	`if optimizer_state_dict is not None and self.restart_training:`
`627`	`628`	`self.optimizer.load_state_dict(optimizer_state_dict)`