From c78b5b2bb74a0be27f4d312ed260b7d3e40060b7 Mon Sep 17 00:00:00 2001 From: Dhruv Gautam <108053562+dhruvji@users.noreply.github.com> Date: Mon, 7 Jul 2025 15:24:24 -0700 Subject: [PATCH] finetuning fix --- src/state/_cli/_tx/_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/state/_cli/_tx/_train.py b/src/state/_cli/_tx/_train.py index b3ce90fd..915dd278 100644 --- a/src/state/_cli/_tx/_train.py +++ b/src/state/_cli/_tx/_train.py @@ -258,7 +258,7 @@ def run_tx_train(cfg: DictConfig): print(f"Loading manual checkpoint from {manual_init}") checkpoint_path = manual_init device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - checkpoint = torch.load(checkpoint_path, map_location=device) + checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) model_state = model.state_dict() checkpoint_state = checkpoint["state_dict"]