Method
__init__
(self, block_type: str, blocks: list[nn.Module], num_blocks: int, blocks_to_swap: int, device: torch.device, debug: bool = False)
Source from the content-addressed store, hash-verified
| 125 | """ |
| 126 | |
| 127 | def __init__(self, block_type: str, blocks: list[nn.Module], num_blocks: int, blocks_to_swap: int, device: torch.device, debug: bool = False): |
| 128 | self.block_type = block_type |
| 129 | self.blocks = blocks |
| 130 | self.num_blocks = num_blocks |
| 131 | self.blocks_to_swap = blocks_to_swap |
| 132 | self.blocks_to_swap_tmp = None |
| 133 | self.device = device |
| 134 | self.debug = debug |
| 135 | |
| 136 | self.thread_pool = ThreadPoolExecutor(max_workers=1) |
| 137 | self.futures = {} |
| 138 | self.cuda_available = device.type == "cuda" |
| 139 | |
| 140 | def swap_weight_devices(self, block_to_cpu: nn.Module, block_to_cuda: nn.Module): |
| 141 | if self.cuda_available: |
Tested by
no test coverage detected