Fix distillation_loss arg typo

Signed-off-by: Austin Liu <[email protected]>
linkedin · Dec 5, 2024 · e87803f · e87803f
1 parent 3a9f125
commit e87803f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/utils.py b/test/utils.py
@@ -525,7 +525,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