import torch


def batch_gather_3_1(inputs, dim):
    """
    Args:
        inputs (batchsz, sen_len, embed_dim)
        dim (batchsz)
    Returns:
        output (batch, embed_dim)
    """
    a = torch.arange(dim.shape[0])
    b = dim.view(-1) - 1
    output = inputs[a, b, :]
    return output