frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix type checking for persistent loads in the weights-only unpickler (#161661)

Browse Source 
The error message here implies that we can only call `self.persistent_load(...)` for ints or tuples, but due to the second part of the type check being inverted, weights-only unpickler will throw an exception iff `pid` is an int.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/161661
Approved by: https://github.com/Skylion007
This commit is contained in:
Ivan Komarov
2025-09-01 19:57:15 +00:00
committed by PyTorch MergeBot
parent d232a95d4a
commit cbfb005f7c
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
torch/_weights_only_unpickler.py
View File

@ -520,7 +520,7 @@ class Unpickler:
elif key[0] == BINPERSID[0]: elif key[0] == BINPERSID[0]:
pid = self.stack.pop() pid = self.stack.pop()
# Only allow persistent load of storage # Only allow persistent load of storage
if type(pid) is not tuple and not type(pid) is not int: if type(pid) is not tuple and type(pid) is not int:
raise UnpicklingError( raise UnpicklingError(
f"persistent_load id must be tuple or int, but got {type(pid)}" f"persistent_load id must be tuple or int, but got {type(pid)}"
) )