def METHOD_NAME(self, emissions, masks=None, beam=None):
    # HACK: we use a beam of tokens to approximate the normalizing factor (which is bad?)
    beam = beam if beam is not None else self.beam
    batch_size, seq_len = emissions.size()[:2]
    beam_emission_scores, beam_targets = emissions.topk(beam, 2)
    beam_transition_score1 = self.E1(beam_targets[:, :-1])  # B x (T-1) x K x D
    beam_transition_score2 = self.E2(beam_targets[:, 1:])  # B x (T-1) x K x D
    beam_transition_matrix = torch.bmm(
        beam_transition_score1.view(-1, beam, self.rank),
        beam_transition_score2.view(-1, beam, self.rank).transpose(1, 2),
    )
    beam_transition_matrix = beam_transition_matrix.view(batch_size, -1, beam, beam)
    traj_tokens, traj_scores = [], []
    finalized_tokens, finalized_scores = [], []
    # compute the normalizer in the log-space
    score = beam_emission_scores[:, 0]  # B x K
    dummy = (
        torch.arange(beam, device=score.device).expand(*score.size()).contiguous()
    )
    for i in range(1, seq_len):
        traj_scores.append(score)
        _score = score[:, :, None] + beam_transition_matrix[:, i - 1]
        _score, _index = _score.max(dim=1)
        _score = _score + beam_emission_scores[:, i]
        if masks is not None:
            score = torch.where(masks[:, i : i + 1], _score, score)
            index = torch.where(masks[:, i : i + 1], _index, dummy)
        else:
            score, index = _score, _index
        traj_tokens.append(index)
    # now running the back-tracing and find the best
    best_score, best_index = score.max(dim=1)
    finalized_tokens.append(best_index[:, None])
    finalized_scores.append(best_score[:, None])
    for idx, scs in zip(reversed(traj_tokens), reversed(traj_scores)):
        previous_index = finalized_tokens[-1]
        finalized_tokens.append(idx.gather(1, previous_index))
        finalized_scores.append(scs.gather(1, previous_index))
    finalized_tokens.reverse()
    finalized_tokens = torch.cat(finalized_tokens, 1)
    finalized_tokens = beam_targets.gather(2, finalized_tokens[:, :, None])[:, :, 0]
    finalized_scores.reverse()
    finalized_scores = torch.cat(finalized_scores, 1)
    finalized_scores[:, 1:] = finalized_scores[:, 1:] - finalized_scores[:, :-1]
    return finalized_scores, finalized_tokens