def METHOD_NAME(self, model_input: Dict[str, torch.Tensor]) -> ReplaceLMOutput:
    enc_prefix, dec_prefix = 'enc_', 'dec_'
    encoder_inputs = slice_batch_dict(model_input, enc_prefix)
    decoder_inputs = slice_batch_dict(model_input, dec_prefix)
    labels = model_input['labels']
    enc_sampled_input_ids, g_mlm_loss = self._replace_tokens(encoder_inputs)
    if self.args.rlm_freeze_generator:
        g_mlm_loss = torch.tensor(0, dtype=torch.float, device=g_mlm_loss.device)
    dec_sampled_input_ids, _ = self._replace_tokens(decoder_inputs, no_grad=True)
    encoder_inputs['input_ids'] = enc_sampled_input_ids
    decoder_inputs['input_ids'] = dec_sampled_input_ids
    # use the un-masked version of labels
    encoder_inputs['labels'] = labels
    decoder_inputs['labels'] = labels
    is_replaced = (encoder_inputs['input_ids'] != labels) & (labels >= 0)
    replace_cnt = is_replaced.long().sum().item()
    total_cnt = (encoder_inputs['attention_mask'] == 1).long().sum().item()
    replace_ratio = torch.tensor(replace_cnt / total_cnt, device=g_mlm_loss.device)
    encoder_out: MaskedLMOutput = self.encoder(
        **encoder_inputs,
        output_hidden_states=True,
        return_dict=True)
    # batch_size x 1 x hidden_dim
    cls_hidden = encoder_out.hidden_states[-1][:, :1]
    # batch_size x seq_length x embed_dim
    dec_inputs_embeds = self.encoder.bert.embeddings(decoder_inputs['input_ids'])
    hiddens = torch.cat([cls_hidden, dec_inputs_embeds[:, 1:]], dim=1)
    attention_mask = self.encoder.get_extended_attention_mask(
        encoder_inputs['attention_mask'],
        encoder_inputs['attention_mask'].shape,
        encoder_inputs['attention_mask'].device
    )
    for layer in self.decoder:
        layer_out = layer(hiddens, attention_mask)
        hiddens = layer_out[0]
    decoder_mlm_loss = self.mlm_loss(hiddens, labels)
    loss = decoder_mlm_loss + encoder_out.loss + g_mlm_loss * self.args.rlm_generator_mlm_weight
    return ReplaceLMOutput(loss=loss,
                           encoder_mlm_loss=encoder_out.loss.detach(),
                           decoder_mlm_loss=decoder_mlm_loss.detach(),
                           g_mlm_loss=g_mlm_loss.detach(),
                           replace_ratio=replace_ratio)