docs: def pred_to_prob(Y_h, k):
      """Converts a 1D tensor of predicted labels into a 2D tensor of probabilistic labels
  
      Args:
          Y_h: an [n], or [n,1] tensor of predicted (int) labels in {1,...,k}
          k: the largest possible label in Y_h
      Returns:
          Y_s: a torch.FloatTensor of shape [n, k] where Y_s[i, j-1] is the probabilistic
              label for item i and label j
      """
      Y_h = Y_h.clone()
      if Y_h.dim() > 1:
          Y_h = Y_h.squeeze()
      assert Y_h.dim() == 1
      assert (Y_h >= 1).all()
      assert (Y_h <= k).all()
      n = Y_h.shape[0]
      Y_s = torch.zeros((n, k), dtype=Y_h.dtype, device=Y_h.device)
      for i, j in enumerate(Y_h):
          Y_s[i, j - 1] = 1.0
      return Y_s