Ensure checkpoint loading works with PyTorch 2.6#14
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| chk_filename = os.path.join(args.checkpoint, args.resume if args.resume else args.evaluate) | ||
| # chk_filename = args.resume or args.evaluate | ||
| print('Loading checkpoint', chk_filename) | ||
| checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage) | ||
| checkpoint = torch.load( | ||
| chk_filename, | ||
| map_location=lambda storage, loc: storage, | ||
| weights_only=False, | ||
| ) |
There was a problem hiding this comment.
Guard torch.load weights_only for pre-2.6 PyTorch
The new torch.load(..., weights_only=False) calls assume the argument exists, but earlier PyTorch releases (e.g. 1.x) do not accept weights_only and raise TypeError: torch.load() got an unexpected keyword argument 'weights_only' before any checkpoint is loaded. The README still advertises support for PyTorch ≥0.4, so this change will crash on previously supported environments. Consider only passing weights_only when running on PyTorch versions that provide the parameter (e.g. via feature detection or try/except). The same pattern is repeated in the other files modified in this commit.
Useful? React with 👍 / 👎.
Summary
Testing
https://chatgpt.com/codex/tasks/task_e_68fdc5efc2ac832799c070a431cf71cf