Add weights_only=True to torch.load #3012
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/3012
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit cac47b8 with merge base 4c1dd9d ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| @@ -34,7 +34,7 @@ | |||
| # using ``load_state_dict()`` method. | |||
|
|
|||
| model = models.vgg16() # we do not specify ``weights``, i.e. create untrained model | |||
| model.load_state_dict(torch.load('model_weights.pth')) | |||
| model.load_state_dict(torch.load('model_weights.pth', weights_only=True)) | |||
There was a problem hiding this comment.
nit: since this is the main landing page for save load, a comment along these lines would be great
"Setting weights_only=True restricts the functions that can be executed during unpickling to those required to load weights, using this flag is a best practice for loading weights."
| @@ -52,7 +52,7 @@ | |||
| ######################## | |||
| # We can then load the model like this: | |||
|
|
|||
| model = torch.load('model.pth') | |||
| model = torch.load('model.pth', weights_only=False), | |||
There was a problem hiding this comment.
nit: I would add a comment here that as noted in https://pytorch.org/docs/main/notes/serialization.html#saving-and-loading-torch-nn-modules, saving state_dicts is the best practice and that we use weights_only=False here because this is loading a model, which is a legacy use case for torch.save
|
Thanks, @mikaylagawarecki! I added your notes - can you take a look? |
* Add weights_only=True to torch.load
* Add weights_only=True to torch.load
Fixes #3007
Fixing the issue and a bunch of other tutorials that use torch.load. In the generated output we have something like this:
We should fix all those warnings.