repeat onehot in template dimension

EleutherAI · Aug 23, 2023 · 777912e · 777912e
1 parent 707ce8e
commit 777912e
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/elk/training/train.py b/elk/training/train.py
@@ -82,10 +82,8 @@ def apply_to_layer(
 
  reporter = CcsReporter(self.net, d, device=device, num_variants=v)
  train_loss = reporter.fit(first_train_h)
-
  labels = to_one_hot(train_gt, k)
- labels = repeat(train_gt, "n -> n v k", v=v, k=k)
- reporter.platt_scale(labels, first_train_h)
+ labels = repeat(labels, "n k -> n v k", v=v)
 
  elif isinstance(self.net, EigenFitterConfig):
  fitter = EigenFitter(