DLR-RM / DLR-RM/stable-baselines3

`torch.load` without `weights_only` parameter is unsafe

Open
#1,852 18 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.