| 12 | |
| 13 | |
| 14 | class OffloadSavedTensorHook: |
| 15 | mode = "release" # ["release", "debug"] |
| 16 | offload_device = "cpu" # cpu or cuda |
| 17 | |
| 18 | @staticmethod |
| 19 | def unpack(packed): |
| 20 | origin_device, x = packed |
| 21 | return x.to(origin_device) |
| 22 | |
| 23 | @staticmethod |
| 24 | def pack(x: torch.Tensor): |
| 25 | if offload_condition(x): |
| 26 | return x.device, x.to(OffloadSavedTensorHook.offload_device) |
| 27 | else: |
| 28 | return x.device, x |
nothing calls this directly
no outgoing calls
no test coverage detected