| 309 | |
| 310 | |
| 311 | def serialize_torch_artifact(artifact: Dict[str, Any]) -> bytes: |
| 312 | assert ( |
| 313 | FakeTensor not in copyreg.dispatch_table |
| 314 | ), "Refusing to stomp on existing FakeTensor reducer" |
| 315 | try: |
| 316 | copyreg.pickle(FakeTensor, _reduce_fake_tensor) |
| 317 | buffer = io.BytesIO() |
| 318 | # This is a workaround for backend's tensor deserialization problem: |
| 319 | # unpickleTensor() always create a tensor on the device where it was originally saved |
| 320 | # This behavior is bad for multi-gpu training, as we wish to directly load the tensor |
| 321 | # on the designated device. |
| 322 | # For now, we simply move the tensor to cpu before saving. |
| 323 | # TODO: this should be fixed by deserialization instead. |
| 324 | torch.save(artifact, buffer) |
| 325 | return buffer.getvalue() |
| 326 | finally: |
| 327 | del copyreg.dispatch_table[FakeTensor] |
| 328 | |
| 329 | |
| 330 | def deserialize_torch_artifact( |