Update train.py

train.py

@@ -5,6 +5,7 @@ import json
 import os
 import random
 import re
+import torch
 
 class ModelTrainer:
     def __init__(self, model_id, system_prompts_path):
@@ -14,16 +15,19 @@ class ModelTrainer:
         with open(system_prompts_path, 'r', encoding='utf-8') as f:
             self.system_prompts = json.load(f)
         
-        # 初始化tokenizer和model - 添加trust_remote_code=True
+        # 修改模型初始化参数
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_id,
-            trust_remote_code=True  # 添加此参数
+            trust_remote_code=True
         )
+        
+        # 修改这部分的初始化参数
         self.model = AutoModelForCausalLM.from_pretrained(
             model_id,
-            trust_remote_code=True,  # 添加此参数
-            low_cpu_mem_usage=True,
-            torch_dtype='float32'
+            trust_remote_code=True,
+            torch_dtype=torch.float32,  # 使用 torch.float32 而不是字符串
+            device_map='auto',          # 自动选择设备
+            low_cpu_mem_usage=True
         )
         
         # 使用更轻量的LoRA配置