import torch


def convert_onehot_to_dense(onehot):
    """
    Args:
        onehot: [batch_size, seq_len, vocab_size]
    Returns: [batch_size, seq_len]
    """
    return torch.argmax(onehot, dim=2)