docs: 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)