import torch


def get_nearest_neg_sample(anchor, neg):
    """Returns a negative sample per batch for G generating new triplet.
    Inputs:
        anchor: [batch, embedding_size]
        neg: [batch, num_neg, embedding_size]
    Returns:
        selected_neg: [batch, embedding_size]
    """
    dist = torch.bmm(neg, anchor.unsqueeze(-1)).squeeze(-1)  # [batch, num_neg]
    max_idx = torch.argmax(dist, dim=1)  # [batch]
    selected_neg = torch.stack([neg[i, max_idx[i]] for i in range(neg.size(0))], dim=0)
    return selected_neg