(self, layer_idx: int, device: torch.device)
| 38 | param.data = param.data.to("cpu", non_blocking=True) |
| 39 | |
| 40 | def get_offlaod_layer(self, layer_idx: int, device: torch.device): |
| 41 | # init stream |
| 42 | if not hasattr(self, "prefetch_stream"): |
| 43 | self.prefetch_stream = torch.cuda.Stream() |
| 44 | |
| 45 | # delete previous layer |
| 46 | torch.cuda.current_stream().synchronize() |
| 47 | self.evict_previous_layer(layer_idx) |
| 48 | |
| 49 | # make sure the current layer is ready |
| 50 | torch.cuda.synchronize(self.prefetch_stream) |
| 51 | |
| 52 | # load next layer |
| 53 | self.prefetch_layer((layer_idx + 1) % len(self.layers), device) |
| 54 | |
| 55 | |
| 56 | def forward( |
no test coverage detected