torch.load with weights_only=True should not be flagged
#1224
Labels
enhancement
New feature or request
Milestone
Comments
lukehinds
added a commit
to lukehinds/bandit
that referenced
this issue
Feb 10, 2025
The PyTorch plugin (B614) has been updated to properly handle the weights_only parameter in torch.load calls. When weights_only=True is specified, PyTorch will only deserialize known safe types, making the operation more secure. I also removed torch.save as there is no certain insecure element as such, saving any file or artifact requires consideration of what it is you are saving. Changes: - Update plugin to only check torch.load calls (not torch.save) - Fix weights_only check to handle both string and boolean True values - Remove map_location check as it doesn't affect security - Update example file to demonstrate both safe and unsafe cases - Update plugin documentation to mention weights_only as a safe alternative The plugin now correctly identifies unsafe torch.load calls while allowing safe usage with weights_only=True to pass without warning. Fixes: PyCQA#1224
Merged
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Is your feature request related to a problem? Please describe.
Currently bandit reports
B614:pytorch_load_savefor the following code:However, when
weights_only=True, PyTorch will only deserialize known safe types, so there should be no security issue.Describe the solution you'd like
B614 should only be reported when
weights_onlyis missing or set toFalse.Describe alternatives you've considered
In PyTorch 2.6,
weights_onlyis slated to default toTrue, so once that's out, it may not be necessary to warn whenweights_onlyis missing either.Additional context
N/A
Love this idea? Give it a 👍. We prioritize fulfilling features with the most 👍.
The text was updated successfully, but these errors were encountered: