(self, inputs)
| 517 | |
| 518 | @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) |
| 519 | def forward(self, inputs): |
| 520 | hidden_states, encoder_hidden_states, temb, freqs_cos, freqs_sin, img_seq_len = inputs |
| 521 | |
| 522 | self.offloader.wait_for_block(self.block_idx) |
| 523 | encoder_hidden_states, hidden_states = self.block( |
| 524 | hidden_states=hidden_states, |
| 525 | encoder_hidden_states=encoder_hidden_states, |
| 526 | temb=temb, |
| 527 | image_rotary_emb=(freqs_cos, freqs_sin), |
| 528 | ) |
| 529 | self.offloader.submit_move_blocks_forward(self.block_idx) |
| 530 | |
| 531 | return make_contiguous(hidden_states, encoder_hidden_states, temb, freqs_cos, freqs_sin, img_seq_len) |
| 532 | |
| 533 | |
| 534 | class OutputWrapper(nn.Module): |
nothing calls this directly
no test coverage detected