def test_center__equivalence() -> None:
    """Check that DINOLoss.update_center is equivalent to Center.update.

    TODO(Guarin, 08/24): Remove this test once DINOLoss uses Center internally.
    """
    criterion = DINOLoss(output_dim=32, center_momentum=0.9)
    center = Center(size=(1, 1, 32), momentum=0.9)
    x = torch.rand(2, 32)
    criterion.update_center(teacher_out=x)
    center.update(x=x)
    assert torch.allclose(criterion.center, center.value)
Check that DINOLoss.update_center is equivalent to Center.update.

    TODO(Guarin, 08/24): Remove this test once DINOLoss uses Center internally.
    
tests/loss/test_dino_loss.py