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