Bug fixing (#441)

Author: oleksost

tests/layers/test_lm_head.py

@@ -39,7 +39,7 @@ def _reverse_kl_loss(
             loss_per_sample = torch.nn.functional.kl_div(
                 teacher_log_probs, student_log_probs, reduction="none", log_target=True
             ).sum(dim=-1)
-            loss = (loss_per_sample * loss_mask.flatten()).sum() / loss_mask.sum()
+            loss = (loss_per_sample * loss_mask.flatten()).mean()
     return loss
 
 

tests/layers/test_ssm.py

@@ -19,11 +19,6 @@
     Apriel2GatedDeltaNet = None
     Apriel2Mamba = None
 
-try:
-    from fast_llm_external_models.apriel_hybrid_ssm.modeling_apriel_hybrid_ssm import KimiDeltaAttention
-except ImportError:
-    KimiDeltaAttention = None
-
 HIDDEN_SIZE = 16
 SEQ_LEN = 65