(self, layer, block_idx, offloader)
| 375 | |
| 376 | class SingleTransformerLayer(nn.Module): |
| 377 | def __init__(self, layer, block_idx, offloader): |
| 378 | super().__init__() |
| 379 | self.layer = layer |
| 380 | self.block_idx = block_idx |
| 381 | self.offloader = offloader |
| 382 | |
| 383 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 384 | def forward(self, inputs): |