Skip to content

Commit

Permalink
Update unwrap_and_save_reload_schedule to use weights_only=False (h…
Browse files Browse the repository at this point in the history
…uggingface#35952)

* fix

* Fix

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
  • Loading branch information
ydshieh and ydshieh authored Jan 29, 2025
1 parent 42c8ccf commit c600e89
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion tests/optimization/test_optimization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit c600e89

Please sign in to comment.