Convert a tensor of indices of any shape `(N, ...)` to a
    tensor of one-hot indicators of shape `(N, num_classes, ...)`.

def to_onehot(indices, num_classes):
    """Convert a tensor of indices of any shape `(N, ...)` to a
    tensor of one-hot indicators of shape `(N, num_classes, ...)`.
    """
    onehot = torch.zeros(indices.shape[0], num_classes, *indices.shape[1:], device=indices.device)
    return onehot.scatter_(1, indices.unsqueeze(1), 1)