add no grad in tests

linkedin · Dec 7, 2024 · 3561525 · 3561525
1 parent 8aa842a
commit 3561525
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/test/utils.py b/test/utils.py
@@ -567,9 +567,10 @@ def concatenated_forward(
         if student_bias is not None:
             student_outputs = student_outputs + student_bias
 
-        teacher_outputs = teacher_input_reshaped @ teacher_weight.t()
-        if teacher_bias is not None:
-            teacher_outputs = teacher_outputs + teacher_bias
+        with torch.no_grad():
+            teacher_outputs = teacher_input_reshaped @ teacher_weight.t()
+            if teacher_bias is not None:
+                teacher_outputs = teacher_outputs + teacher_bias
 
         student_logits = student_outputs.view(student_batch_seq_len_size, -1).float()
         teacher_logits = teacher_outputs.view(teacher_batch_seq_len_size, -1).float()