saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix test_save_load_transform. (#140494)

Browse source 
test test_save_load_transform in [test_transforms.py](https://github.com/pytorch/pytorch/blob/main/test/distributions/test_transforms.py)

_pytest test_transforms.py -k test_save_load_transform_

error message:

```
.
.
.
  File "/workspace/pytorch/test/distributions/test_transforms.py", line 555, in test_save_load_transform
    other = torch.load(stream)
            ^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/serialization.py", line 1444, in load
    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
_pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
	(1) Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
	(2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.
	WeightsUnpickler error: Unsupported global: GLOBAL torch.distributions.transformed_distribution.TransformedDistribution was not an allowed global by default. Please use `torch.serialization.add_safe_globals([TransformedDistribution])` or the `torch.serialization.safe_globals([TransformedDistribution])` context manager to allowlist this global if you trust this class/function.

```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/140494
Approved by: https://github.com/mikaylagawarecki
This commit is contained in:
sifengyang 2024-11-19 04:36:03 +00:00 committed by PyTorch MergeBot
parent d472a5f680
commit 51d4338716
1 changed files with 4 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
test/distributions/test_transforms.py
View file

@ -552,7 +552,10 @@ def test_save_load_transform():
stream = io.BytesIO()
torch.save(dist, stream)
stream.seek(0)
other = torch.load(stream)
with torch.serialization.safe_globals(
[TransformedDistribution, AffineTransform, Normal]
):
other = torch.load(stream)
assert torch.allclose(log_prob, other.log_prob(x))