diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py
index 2610d1a..c5ee3b2 100644
--- a/trainer/train_full_sft.py
+++ b/trainer/train_full_sft.py
@@ -128,7 +128,7 @@ if __name__ == "__main__":
     train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
-    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=('cuda' in args.device))
     
     # ========== 6. 从ckp恢复状态 ==========
     start_epoch, start_step = 0, 0
diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py
index a5d2fa2..7ac0d2b 100644
--- a/trainer/train_pretrain.py
+++ b/trainer/train_pretrain.py
@@ -127,7 +127,7 @@ if __name__ == "__main__":
     train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
-    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=(device_type == 'cuda'))
     
     # ========== 6. 从ckp恢复状态 ==========
     start_epoch, start_step = 0, 0