Update train.py

train.py

@@ -9,6 +9,9 @@ import torch
 
 class ModelTrainer:
     def __init__(self, model_id, system_prompts_path):
+        # 确保临时文件夹存在
+        os.makedirs("temp_model_dir", exist_ok=True)
+        
         self.model_id = model_id
         
         # 加载系统提示词
@@ -27,7 +30,9 @@ class ModelTrainer:
             trust_remote_code=True,
             torch_dtype=torch.float32,  # 使用 torch.float32 而不是字符串
             device_map='auto',          # 自动选择设备
-            low_cpu_mem_usage=True
+            low_cpu_mem_usage=True,
+            offload_folder="temp_model_dir",  # 添加临时文件夹
+            use_safetensors=True  # 使用 safetensors
         )
         
         # 使用更轻量的LoRA配置