ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #20 from ccx1324/lora-device-fix

Browse Source 
fix: LoRA device mismatch and checkpoint resume
This commit is contained in:
ViperEkura
2026-07-20 19:38:10 +08:00
committed by GitHub
parent 958df58f9d 5a942527b2
commit a6e920fdb0
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/trainer/train_context.py
+2 -2
View File
@@ -110,8 +110,6 @@ class TrainContextBuilder:
def _before_wrap(m): def _before_wrap(m):
m = m.to(device=device) m = m.to(device=device)
if preloaded_state_dict is not None:
m.load_state_dict(preloaded_state_dict, strict=False)
if cfg.lora is not None: if cfg.lora is not None:
inject_lora( inject_lora(
m, m,
@@ -119,6 +117,8 @@ class TrainContextBuilder:
alpha=cfg.lora.alpha, alpha=cfg.lora.alpha,
target_modules=set(cfg.lora.target_modules), target_modules=set(cfg.lora.target_modules),
) )
if preloaded_state_dict is not None:
m.load_state_dict(preloaded_state_dict, strict=False)
return m return m
context = TrainContext( context = TrainContext(