| 154 | return self.decoder(input_ids, position_ids, attention_mask, encoder_outputs=encoder_outputs, ids_restore=ids_restore, **kw_args) |
| 155 | |
| 156 | def forward(self, input_ids, enc_position_ids, dec_position_ids, *, enc_attention_mask=None, dec_attention_mask=None, **kw_args): |
| 157 | if enc_attention_mask is None: |
| 158 | enc_attention_mask = torch.ones(1, 1, dtype=self.encoder.transformer.word_embeddings.weight.dtype, device=input_ids.device) |
| 159 | encoder_outputs, *encoder_mems = self.encode(input_ids, enc_position_ids, enc_attention_mask, **kw_args) |
| 160 | decoder_outputs, *decoder_mems = self.decode(input_ids, dec_position_ids, dec_attention_mask, encoder_outputs=encoder_outputs, ids_restore=encoder_mems[0]["ids_restore"], **kw_args) |
| 161 | return encoder_outputs, decoder_outputs, encoder_mems, decoder_mems |
| 162 | |
| 163 | def unpatchify(self, x): |
| 164 | """ |