def METHOD_NAME(model, input_ids, beam_width, alpha):
    '''
        model: the generation model, e.g., gpt2
        input_ids: 1 x seqlen
    '''
    device = next(model.parameters()).device
    with torch.no_grad():
        input_ids = torch.as_tensor(input_ids, device=device)
        if input_ids.ndim == 1:
            input_ids = input_ids.view(1, -1)
        outputs = model(**{"input_ids": input_ids, "output_hidden_states": True})
        logits = outputs["logits"]
        prev_hidden_states = outputs["logits"][-1]
    _, _, vocab_size = logits.size()
    seqlen, embed_dim = prev_hidden_states.size()
    p = random.uniform(0, 1)
    logit_for_next_step = logits[:,-1,:]
    assert logit_for_next_step.size() == torch.Size([1, vocab_size])
    next_probs = F.softmax(logit_for_next_step, dim = -1)
    assert next_probs.size() == logit_for_next_step.size()
    _, top_k_ids = torch.topk(logit_for_next_step, dim = -1, k = beam_width)
    assert top_k_ids.size() == torch.Size([1, beam_width])
    top_k_probs = torch.gather(next_probs, dim = 1, index=top_k_ids)
    assert top_k_probs.size() == top_k_ids.size()
    # compute new hidden 
    expanded_context = [input_ids for _ in range(beam_width)]
    expanded_context = torch.cat(expanded_context, dim = 0)
    assert expanded_context.size() == torch.Size([beam_width, seqlen])
    top_k_ids = top_k_ids.view(beam_width, 1)
    next_input_ids = torch.cat([expanded_context, top_k_ids], dim = -1)
    assert next_input_ids.size() == torch.Size([beam_width, seqlen+1])
    #new_hidden_states, next_logits = model.compute_logits_and_hidden_states(next_input_ids)
    with torch.no_grad():
        #next_input_ids = torch.tensor(next_input_ids, device=device)
        if next_input_ids.ndim == 1:
            next_input_ids = next_input_ids.view(1, -1)
        outputs = model(**{"input_ids": next_input_ids, "output_hidden_states": True})
        next_logits = outputs["logits"]
        new_hidden_states = outputs["logits"]
    assert new_hidden_states.size() == torch.Size([beam_width, seqlen+1, embed_dim])
    context_hidden = new_hidden_states[:,:seqlen,:]
    assert context_hidden.size() == torch.Size([beam_width, seqlen, embed_dim])
    next_hidden = new_hidden_states[:,seqlen:,:]
    assert next_hidden.size() == torch.Size([beam_width, 1, embed_dim])
    next_id = ranking(context_hidden, next_hidden, top_k_ids, top_k_probs, alpha)       
    next_input_ids = torch.cat([input_ids, next_id], dim = -1)
    assert next_input_ids.size() == torch.Size([1, seqlen+1])
    return next_input_ids