Fix distillation_loss arg typo

Signed-off-by: Austin Liu <[email protected]>
linkedin · Dec 7, 2024 · 7acb5ca · 7acb5ca
1 parent 5662554
commit 7acb5ca
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/utils.py b/test/utils.py
@@ -536,7 +536,7 @@ def __init__(
         self.temperature = temperature
 
     @abstractmethod
-    def distillation_loss(self, student_logps, teacher_logps):
+    def distillation_loss(self, student_logits, teacher_logits):
         """Abstract method for computing distillation loss."""
         pass