def METHOD_NAME(logits, tau, just_prob=False, num=2, eps=1e-7):
    if tau <= 0:
        new_logits = logits
        probs = nn.functional.softmax(new_logits, dim=1)
    else:
        while True:  # a trick to avoid the gumbels bug
            gumbels = -torch.empty_like(logits).exponential_().log()
            new_logits = (logits.log_softmax(dim=1) + gumbels) / tau
            probs = nn.functional.softmax(new_logits, dim=1)
            if (
                (not torch.isinf(gumbels).any())
                and (not torch.isinf(probs).any())
                and (not torch.isnan(probs).any())
            ):
                break
    if just_prob:
        return probs
    # with torch.no_grad(): # add eps for unexpected torch error
    #  probs = nn.functional.softmax(new_logits, dim=1)
    #  selected_index = torch.multinomial(probs + eps, 2, False)
    with torch.no_grad():  # add eps for unexpected torch error
        probs = probs.cpu()
        selected_index = torch.multinomial(probs + eps, num, False).to(logits.device)
    selected_logit = torch.gather(new_logits, 1, selected_index)
    selcted_probs = nn.functional.softmax(selected_logit, dim=1)
    return selected_index, selcted_probs