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