import torch


def one_hot_embedding(labels, num_classes=2):
    """Embedding labels to one-hot form.
    Args:
      labels: (LongTensor) class labels, sized [B,N].
      num_classes: (int) number of classes.
    Returns:
      (tensor) encoded labels, sized [N,#classes].
    """
    B, N = labels.size()
    # labels = labels.view(-1, 1)  # [B,N]->[B*N,1]
    labels = labels.view(int(B * N), 1)
    y = torch.FloatTensor(labels.size()[0], num_classes)  # [B*N, D]
    y.zero_()
    y.scatter_(1, labels, 1)
    return y