| 265 | |
| 266 | |
| 267 | class SavingProxyForStorage: |
| 268 | def __init__(self, obj, saver, protocol_version=5): |
| 269 | self.protocol_version = protocol_version |
| 270 | self.saver = saver |
| 271 | if not (isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj)): |
| 272 | raise TypeError(f"expected storage, not {type(obj)}") |
| 273 | |
| 274 | # this logic is taken from PyTorch 2.0+ torch/serialization.py |
| 275 | if isinstance(obj, torch.storage.TypedStorage): |
| 276 | # PT upstream wants to deprecate this eventually... |
| 277 | storage = obj._untyped_storage |
| 278 | storage_type_str = obj._pickle_storage_type() |
| 279 | storage_type = getattr(torch, storage_type_str) |
| 280 | storage_numel = obj._size() |
| 281 | else: |
| 282 | storage = obj |
| 283 | storage_type = normalize_storage_type(type(obj)) |
| 284 | storage_numel = storage.nbytes() |
| 285 | |
| 286 | storage_key = saver._write_storage_and_return_key(storage) |
| 287 | location = torch.serialization.location_tag(storage) |
| 288 | |
| 289 | self.storage_info = ("storage", storage_type, storage_key, location, storage_numel) |
| 290 | |
| 291 | def __reduce_ex__(self, protocol_version): |
| 292 | assert False, "this should be handled with out of band" |
| 293 | |
| 294 | |
| 295 | class SavingProxyForTensor: |