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

Minor fix in package exporter (#90306)

Browse Source 
Summary:
As title.
Saw this while working on another diff.
`storage` won't be defined in the `else` case. But this causes pyre to freak out.

Test Plan: Unit tests.

Differential Revision: D41751229

Pull Request resolved: https://github.com/pytorch/pytorch/pull/90306
Approved by: https://github.com/PaliC
This commit is contained in:
Atul Jangra
2022-12-27 18:01:59 +00:00
committed by PyTorch MergeBot
parent 06bdd491fb
commit 467d269ad1
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
torch/package/package_exporter.py
View File

@ -891,22 +891,25 @@ class PackageExporter:
def _persistent_id(self, obj):
if torch.is_storage(obj) or isinstance(obj, torch.storage.TypedStorage):
storage: Storage
if isinstance(obj, torch.storage.TypedStorage):
# TODO: Once we decide to break serialization FC, we can
# remove this case
untyped_storage = obj._untyped_storage
storage_type_str = obj.pickle_storage_type()
storage_type = getattr(torch, storage_type_str)
storage = cast(Storage, untyped_storage)
storage_numel = obj.size()
elif isinstance(obj, torch.UntypedStorage):
untyped_storage = obj
storage = cast(Storage, untyped_storage)
storage_type = normalize_storage_type(type(storage))
storage_numel = storage.nbytes()
else:
raise RuntimeError(f"storage type not recognized: {type(obj)}")
storage: Storage = cast(Storage, untyped_storage)
location = location_tag(storage)
# serialize storage if not already written