From c600e89f5c989959dea7701489d93a6dddeca3bb Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Wed, 29 Jan 2025 14:30:57 +0100 Subject: [PATCH] Update `unwrap_and_save_reload_schedule` to use `weights_only=False` (#35952) * fix * Fix --------- Co-authored-by: ydshieh --- tests/optimization/test_optimization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/optimization/test_optimization.py b/tests/optimization/test_optimization.py index adf3039bb4..6982583d2b 100644 --- a/tests/optimization/test_optimization.py +++ b/tests/optimization/test_optimization.py @@ -59,7 +59,7 @@ def unwrap_and_save_reload_schedule(scheduler, num_steps=10): file_name = os.path.join(tmpdirname, "schedule.bin") torch.save(scheduler.state_dict(), file_name) - state_dict = torch.load(file_name) + state_dict = torch.load(file_name, weights_only=False) scheduler.load_state_dict(state_dict) return lrs