(self, layer, block_idx, offloader)
| 328 | |
| 329 | class DoubleTransformerLayer(nn.Module): |
| 330 | def __init__(self, layer, block_idx, offloader): |
| 331 | super().__init__() |
| 332 | self.layer = layer |
| 333 | self.block_idx = block_idx |
| 334 | self.offloader = offloader |
| 335 | |
| 336 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 337 | def forward(self, inputs): |