def one_hot_embedding(labels, num_classes=10):
 # Convert to One Hot Encoding
 device = get_device()
 y = torch.eye(num_classes).to(device)
 return y[labels]