(self, layer, block_idx, offloader)
| 224 | |
| 225 | class TransformerLayer(nn.Module): |
| 226 | def __init__(self, layer, block_idx, offloader): |
| 227 | super().__init__() |
| 228 | self.layer = layer |
| 229 | self.block_idx = block_idx |
| 230 | self.offloader = offloader |
| 231 | |
| 232 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 233 | def forward(self, inputs): |