DLR-RM / DLR-RM/stable-baselines3
`torch.load` without `weights_only` parameter is unsafe
- Dominant language
- Python
- Stars
- 13.8k
- Forks
- 2.2k
- Avg merge
- 1h 35m
- Merged PRs (30d)
- 2
Description
This is found via https://github.com/pytorch-labs/torchfix/
`torch.load` without `weights_only` parameter is unsafe. Explicitly set `weights_only` to False only if you trust the data you load and full pickle functionality is needed, otherwise set `weights_only=True`.
stable_baselines3/common/policies.py:176:27
```
--- /home/sdym/repos/stable-baselines3/stable_baselines3/common/policies.py
+++ /home/sdym/repos/stable-baselines3/stable_baselines3/common/policies.py
@@ -171,11 +171,11 @@
:param path:
:param device: Device on which the policy should be loaded.
:return:
"""
device = get_device(device)
- saved_variables = th.load(path, map_location=device)
+ saved_variables = th.load(path, map_location=device, weights_only=True)
# Create policy object
model = cls(**saved_variables["data"])
# Load weights
model.load_state_dict(saved_variables["state_dict"])
```
stable_baselines3/common/save_util.py:450:33
```
--- /home/sdym/repos/stable-baselines3/stable_baselines3/common/save_util.py
+++ /home/sdym/repos/stable-baselines3/stable_baselines3/common/save_util.py
@@ -445,11 +445,11 @@
file_content.write(param_file.read())
# go to start of file
file_content.seek(0)
# Load the parameters with the right ``map_location``.
# Remove ".pth" ending with splitext
- th_object = th.load(file_content, map_location=device)
+ th_object = th.load(file_content, map_location=device, weights_only=True)
# "tensors.pth" was renamed "pytorch_variables.pth" in v0.9.0, see PR #138
if file_path == "pytorch_variables.pth" or file_path == "tensors.pth":
# PyTorch variables (not state_dicts)
pytorch_variables = th_object
else:
```
Contributor guide
Research direction
Review the torch.load calls in stable_baselines3/common/policies.py and stable_baselines3/common/save_util.py, starting with the shown locations and their surrounding loading logic. Confirm both calls explicitly set weights_only and verify that policy and parameter loading still works for the supported saved-data cases.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, security
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100