Merge nucleic/sleek-ember-seal-uady into dev

tests/test_train.py

@@ -12,6 +12,19 @@ import train
 
 
 class MetricsTests(unittest.TestCase):
+    def test_distillation_loss_is_zero_for_matching_logits_and_backpropagates(self):
+        teacher = torch.tensor([[2.0, 0.0, -1.0]])
+        student = teacher.clone().requires_grad_(True)
+        loss = train.knowledge_distillation_loss(
+            torch,
+            student,
+            teacher,
+            temperature=2.0,
+        ).mean()
+        self.assertAlmostEqual(0.0, loss.item(), places=6)
+        loss.backward()
+        self.assertIsNotNone(student.grad)
+
     def test_qat_replacements_keep_checkpoint_keys_and_gradients(self):
         model = torch.nn.Sequential(
             torch.nn.Embedding(16, 8),