| 192 | |
| 193 | |
| 194 | class LazyLoadingUnpickler(pickle.Unpickler): |
| 195 | def __init__(self, file, zipfile_context): |
| 196 | super().__init__(file) |
| 197 | self.zipfile_context = zipfile_context |
| 198 | |
| 199 | def find_class(self, module, name): |
| 200 | res = super().find_class(module, name) |
| 201 | if module == "torch._utils" and name == "_rebuild_tensor_v2": |
| 202 | return partial(NotYetLoadedTensor.rebuild_tensor_v2, archiveinfo=self) |
| 203 | if module == "torch._tensor" and name == "_rebuild_from_type_v2": |
| 204 | return partial(NotYetLoadedTensor.rebuild_from_type_v2, archiveinfo=self) |
| 205 | if module == "torch._utils" and name == "_rebuild_parameter": |
| 206 | return partial(NotYetLoadedTensor.rebuild_parameter, archiveinfo=self) |
| 207 | return res |
| 208 | |
| 209 | def persistent_load(self, pid): |
| 210 | name, cls, fn, device, size = pid |
| 211 | with warnings.catch_warnings(): |
| 212 | warnings.simplefilter("ignore") |
| 213 | s = torch.storage.TypedStorage(dtype=cls().dtype, device="meta") |
| 214 | s.archiveinfo = pid |
| 215 | return s |
| 216 | |
| 217 | |
| 218 | class lazy_load: |