Skip to content

Commit 0c51d92

Browse files
authored
fix state dict adapter in forge engine (#1563)
This was broken by #1526, updating to match the changes in train.py. Tested via a run in forge
1 parent 21416c4 commit 0c51d92

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

torchtitan/experiments/forge/engine.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,9 @@ def __init__(self, job_config: ForgeJobConfig):
216216
states={"train_state": self},
217217
checkpoint_config=job_config.checkpoint,
218218
sd_adapter=(
219-
self.train_spec.state_dict_adapter(model_args)
219+
self.train_spec.state_dict_adapter(
220+
model_args, job_config.model.hf_assets_path
221+
)
220222
if self.train_spec.state_dict_adapter
221223
else None
222224
),

0 commit comments

Comments
 (0)