| 138 | |
| 139 | |
| 140 | class SavingProxyForTensor: |
| 141 | def __init__(self, tensor, saver, protocol_version=5): |
| 142 | self.protocol_version = protocol_version |
| 143 | self.reduce_ret_fn, reduce_args = tensor.__reduce_ex__(protocol_version) |
| 144 | if reduce_args[0] == torch._utils._rebuild_tensor_v2: |
| 145 | # for Tensors with Python attributes |
| 146 | (a0, a1, (storage, *a2_other), *other_reduce_args) = reduce_args |
| 147 | assert isinstance( |
| 148 | storage, torch.storage.TypedStorage |
| 149 | ), 'Please check for updates' |
| 150 | storage_proxy = SavingProxyForStorage( |
| 151 | storage, saver, protocol_version=protocol_version |
| 152 | ) |
| 153 | self.reduce_args = (a0, a1, (storage_proxy, *a2_other), *other_reduce_args) |
| 154 | else: |
| 155 | (storage, *other_reduce_args) = reduce_args |
| 156 | assert isinstance( |
| 157 | storage, torch.storage.TypedStorage |
| 158 | ), 'Please check for updates' |
| 159 | storage_proxy = SavingProxyForStorage( |
| 160 | storage, saver, protocol_version=protocol_version |
| 161 | ) |
| 162 | self.reduce_args = (storage_proxy, *other_reduce_args) |
| 163 | |
| 164 | def __reduce_ex__(self, protocol_version): |
| 165 | if protocol_version != self.protocol_version: |
| 166 | raise RuntimeError( |
| 167 | f'Unexpected protocol version: expected {self.protocol_version}, got {protocol_version}' |
| 168 | ) |
| 169 | return self.reduce_ret_fn, self.reduce_args |
| 170 | |
| 171 | |
| 172 | class IncrementalPyTorchPickler(pickle.Pickler): |