import torch


def to_one_hot(labels, d_classes):
    """
    Args:
        labels (Tensor): integer tensor of shape (d_batch, *)
        d_classes (int): number of classes
    Returns:
        (Tensor): float tensor of shape (d_batch, *, d_classes), one hot representation of the labels
    """
    return torch.zeros(*labels.size(), d_classes, device=labels.device).scatter_(-1, labels.unsqueeze(-1), 1)