(self, inputs)
| 231 | |
| 232 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 233 | def forward(self, inputs): |
| 234 | h, attn_mask, adaln_input, sizes, *freqs_cis = inputs |
| 235 | |
| 236 | self.offloader.wait_for_block(self.block_idx) |
| 237 | h = self.layer(h, attn_mask, freqs_cis, adaln_input) |
| 238 | self.offloader.submit_move_blocks_forward(self.block_idx) |
| 239 | |
| 240 | return make_contiguous(h, attn_mask, adaln_input, sizes, *freqs_cis) |
| 241 | |
| 242 | |
| 243 | class FinalLayer(nn.Module): |
nothing calls this directly
no test coverage detected