Moves the previous layer cache to the CPU
(self, layer_idx: int)
| 32 | param.data = param.data.to(device, non_blocking=True) |
| 33 | |
| 34 | def evict_previous_layer(self, layer_idx: int): |
| 35 | "Moves the previous layer cache to the CPU" |
| 36 | prev_layer_idx = layer_idx - 1 |
| 37 | for name, param in self.layers[prev_layer_idx].named_parameters(): |
| 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 |