def decode_entity(x, mask):
    """Decode sequences of entities from weight matrix

    Args:
        x (torch.Tensor): output with shape (B, T, num_entities)
        mask (torch.BoolTensor): (B, T)

    Returns:
        (list[list[int]]): best sequences of entities of this batch, representing in indexes (B, *)
    """
    first_invalid = mask.sum(1)  # (B,)

    preds = x.argmax(dim=-1)  # (B, T)
    path = [preds[i].data[:first_invalid[i].item()].tolist() for i in range(preds.shape[0])]  # (B, *)
    return path