Update app.py

app.py

@@ -1,15 +1,17 @@
+from accelerate import Accelerator
 from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
 from peft import LoraConfig, get_peft_model, TaskType
 from datasets import load_dataset
 import torch
-from accelerate import Accelerator
 
 def main():
-    accelerator = Accelerator()#初始化
+    # 初始化 Accelerator
+    accelerator = Accelerator()
+
     # 基础模型位置
     model_name = "dushuai112233/Qwen2-1.5B-Instruct"
     # 设备
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = accelerator.device
 
     # 加载分词器和模型
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
@@ -54,6 +56,9 @@ def main():
         save_total_limit=2,                # 最大保存模型数
     )
 
+    # 将模型移到设备
+    model.to(device)
+
     # Define the Trainer
     trainer = Trainer(
         model=model,                       # 训练的模型
@@ -61,7 +66,6 @@ def main():
         train_dataset=train_dataset,       # 训练数据集
         eval_dataset=val_dataset,          # 验证数据集
         tokenizer=tokenizer,               # 用于预处理的分词器
-        accelerator=accelerator,
     )
 
     # Start Training