File size: 346 Bytes
46cc63a
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
"""Golden Baseline — R-Drop loss and frozen inference mode."""

from __future__ import annotations

import torch

from src.models.transformer_trainer import _symmetric_kl


def test_symmetric_kl_zero_for_identical_logits():
    logits = torch.tensor([[2.0, -1.0], [0.5, 0.5]])
    kl = _symmetric_kl(logits, logits)
    assert kl.item() < 1e-5