remove extra line for label

EleutherAI · Aug 27, 2023 · 6ed5eca · 6ed5eca
1 parent f678c90
commit 6ed5eca
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/elk/training/train.py b/elk/training/train.py
@@ -82,8 +82,7 @@ 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(labels, "n k -> n v k", v=v)
+ labels = repeat(to_one_hot(train_gt, k), "n k -> n v k", v=v)
  reporter.platt_scale(labels, first_train_h)
 
  elif isinstance(self.net, EigenFitterConfig):