56 lines
2.8 KiB
Diff
56 lines
2.8 KiB
Diff
diff -ur pytorch-v2.5.1/test/test_serialization.py pytorch-v2.5.1-patched/test/test_serialization.py
|
|
--- pytorch-v2.5.1/test/test_serialization.py 2024-10-29 14:00:43.000000000 -0400
|
|
+++ pytorch-v2.5.1-patched/test/test_serialization.py 2025-12-25 00:29:50.653124481 -0500
|
|
@@ -459,7 +459,11 @@
|
|
b += [a[0].storage()]
|
|
b += [a[0].reshape(-1)[1:4].clone().storage()]
|
|
path = download_file('https://download.pytorch.org/test_data/legacy_serialized.pt')
|
|
- c = torch.load(path, weights_only=weights_only)
|
|
+ if weights_only:
|
|
+ with self.assertRaisesRegex(RuntimeError,
|
|
+ "Cannot use ``weights_only=True`` with files saved in the legacy .tar format."):
|
|
+ c = torch.load(path, weights_only=weights_only)
|
|
+ c = torch.load(path, weights_only=False)
|
|
self.assertEqual(b, c, atol=0, rtol=0)
|
|
self.assertTrue(isinstance(c[0], torch.FloatTensor))
|
|
self.assertTrue(isinstance(c[1], torch.FloatTensor))
|
|
diff -ur pytorch-v2.5.1/torch/serialization.py pytorch-v2.5.1-patched/torch/serialization.py
|
|
--- pytorch-v2.5.1/torch/serialization.py 2024-10-29 14:00:48.000000000 -0400
|
|
+++ pytorch-v2.5.1-patched/torch/serialization.py 2025-12-25 00:33:44.188969304 -0500
|
|
@@ -83,6 +83,12 @@
|
|
|
|
IS_WINDOWS = sys.platform == "win32"
|
|
|
|
+UNSAFE_MESSAGE = (
|
|
+ "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."
|
|
+)
|
|
+
|
|
if not IS_WINDOWS:
|
|
from mmap import MAP_PRIVATE, MAP_SHARED
|
|
else:
|
|
@@ -1228,11 +1234,6 @@
|
|
>>> torch.load("module.pt", encoding="ascii", weights_only=False)
|
|
"""
|
|
torch._C._log_api_usage_once("torch.load")
|
|
- UNSAFE_MESSAGE = (
|
|
- "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."
|
|
- )
|
|
DOCS_MESSAGE = (
|
|
"\n\nCheck the documentation of torch.load to learn more about types accepted by default with "
|
|
"weights_only https://pytorch.org/docs/stable/generated/torch.load.html."
|
|
@@ -1482,6 +1483,11 @@
|
|
with closing(
|
|
tarfile.open(fileobj=f, mode="r:", format=tarfile.PAX_FORMAT)
|
|
) as tar, mkdtemp() as tmpdir:
|
|
+ if pickle_module is _weights_only_unpickler:
|
|
+ raise RuntimeError(
|
|
+ "Cannot use ``weights_only=True`` with files saved in the "
|
|
+ "legacy .tar format. " + UNSAFE_MESSAGE
|
|
+ )
|
|
tar.extract("storages", path=tmpdir)
|
|
with open(os.path.join(tmpdir, "storages"), "rb", 0) as f:
|
|
num_storages = pickle_module.load(f, **pickle_load_args)
|