| 10 | |
| 11 | class AutoWrappedModule(torch.nn.Module): |
| 12 | def __init__(self, module: torch.nn.Module, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device): |
| 13 | super().__init__() |
| 14 | self.module = module.to(dtype=offload_dtype, device=offload_device) |
| 15 | self.offload_dtype = offload_dtype |
| 16 | self.offload_device = offload_device |
| 17 | self.onload_dtype = onload_dtype |
| 18 | self.onload_device = onload_device |
| 19 | self.computation_dtype = computation_dtype |
| 20 | self.computation_device = computation_device |
| 21 | self.state = 0 |
| 22 | |
| 23 | def offload(self): |
| 24 | if self.state == 1 and (self.offload_dtype != self.onload_dtype or self.offload_device != self.onload_device): |