Method
__init__
(
self,
module: torch.nn.Module,
offload_dtype,
offload_device,
onload_dtype,
onload_device,
computation_dtype,
computation_device,
)
Source from the content-addressed store, hash-verified
| 20 | |
| 21 | class AutoWrappedModule(torch.nn.Module): |
| 22 | def __init__( |
| 23 | self, |
| 24 | module: torch.nn.Module, |
| 25 | offload_dtype, |
| 26 | offload_device, |
| 27 | onload_dtype, |
| 28 | onload_device, |
| 29 | computation_dtype, |
| 30 | computation_device, |
| 31 | ): |
| 32 | super().__init__() |
| 33 | self.module = module.to(dtype=offload_dtype, device=offload_device) |
| 34 | self.offload_dtype = offload_dtype |
| 35 | self.offload_device = offload_device |
| 36 | self.onload_dtype = onload_dtype |
| 37 | self.onload_device = onload_device |
| 38 | self.computation_dtype = computation_dtype |
| 39 | self.computation_device = computation_device |
| 40 | self.state = 0 |
| 41 | |
| 42 | def offload(self): |
| 43 | if self.state == 1 and ( |
Tested by
no test coverage detected