From 3cc6d0d074e3b020e183b2c8ef6703677989ec19 Mon Sep 17 00:00:00 2001
From: Lycoris <149505619+math-zhuxy@users.noreply.github.com>
Date: Mon, 26 May 2025 18:30:46 +0800
Subject: [PATCH] Update train_lora.py

---
 trainer/train_lora.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/trainer/train_lora.py b/trainer/train_lora.py
index 3a799f0..096a63c 100644
--- a/trainer/train_lora.py
+++ b/trainer/train_lora.py
@@ -19,6 +19,7 @@ from transformers import AutoTokenizer
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from dataset.lm_dataset import SFTDataset
 from model.model_lora import load_lora, save_lora, apply_lora
+import safesentors.torch
 
 warnings.filterwarnings('ignore')
 
@@ -99,7 +100,8 @@ def init_model(lm_config):
     model = MiniMindForCausalLM(lm_config)
     if args.minimind2:
         model_data_path = os.path.join(current_dir, '..', 'MiniMind2')
-        model.from_pretrained(model_data_path)
+        state_dict = safetensors.torch.load_file("model.safetensors", device=args.device)
+        model.load_state_dict(state_dict, strict= True)
         return model.to(args.device), tokenizer
     moe_path = '_moe' if lm_config.use_moe else ''
     ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'