def METHOD_NAME(self, positive_batch, batch_filter):
    # shape: (batch_size, 1, num_neg)
    positive_batch = positive_batch.unsqueeze(dim=1)
    if batch_filter is not None:
        positive_batch = positive_batch[batch_filter]
    return positive_batch