def METHOD_NAME(
    self,
    src_tokens,
    src_lengths=None,  # unused
    incremental_state: Optional[Dict[str, List[torch.Tensor]]] = None,
    encoder_out=None,
):
    if incremental_state is not None:  # used during inference
        mems = self.get_incremental_state(incremental_state, "mems")
        src_tokens = src_tokens[:, -1:]  # only keep the most recent token
    else:
        mems = self._mems
    output = self.model(
        input_ids=src_tokens,
        mems=mems,
        return_dict=False,
    )
    if len(output) >= 2:
        if incremental_state is not None:
            self.set_incremental_state(incremental_state, "mems", output[1])
        else:
            self._mems = output[1]
    return (output[0],)