| 104 | |
| 105 | |
| 106 | class SavingProxyForStorage: |
| 107 | def __init__(self, obj, saver, protocol_version=5): |
| 108 | self.protocol_version = protocol_version |
| 109 | self.saver = saver |
| 110 | if not (isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj)): |
| 111 | raise TypeError(f'expected storage, not {type(obj)}') |
| 112 | |
| 113 | # this logic is taken from PyTorch 2.0+ torch/serialization.py |
| 114 | if isinstance(obj, torch.storage.TypedStorage): |
| 115 | # PT upstream wants to deprecate this eventually... |
| 116 | storage = obj._untyped_storage |
| 117 | storage_type_str = obj._pickle_storage_type() |
| 118 | storage_type = getattr(torch, storage_type_str) |
| 119 | storage_numel = obj._size() |
| 120 | else: |
| 121 | storage = obj |
| 122 | storage_type = normalize_storage_type(type(obj)) |
| 123 | storage_numel = storage.nbytes() |
| 124 | |
| 125 | storage_key = saver._write_storage_and_return_key(storage) |
| 126 | location = torch.serialization.location_tag(storage) |
| 127 | |
| 128 | self.storage_info = ( |
| 129 | 'storage', |
| 130 | storage_type, |
| 131 | storage_key, |
| 132 | location, |
| 133 | storage_numel, |
| 134 | ) |
| 135 | |
| 136 | def __reduce_ex__(self, protocol_version): |
| 137 | assert False, 'this should be handled with out of band' |
| 138 | |
| 139 | |
| 140 | class SavingProxyForTensor: |