def _assert_dino_loss_equal_to_original(
    batch_size: int = 3,
    n_global: int = 2,  # number of global views
    n_local: int = 6,  # number of local views
    output_dim: int = 4,
    warmup_teacher_temp: float = 0.04,
    teacher_temp: float = 0.04,
    warmup_teacher_temp_epochs: int = 30,
    student_temp: float = 0.1,
    center_momentum: float = 0.9,
    epoch: int = 0,
    n_epochs: int = 100,
) -> None:
    """Runs test with the given input parameters."""
    loss_fn = DINOLoss(
        output_dim=output_dim,
        teacher_temp=teacher_temp,
        student_temp=student_temp,
        center_momentum=center_momentum,
    )

    orig_loss_fn = OriginalDINOLoss(
        out_dim=output_dim,
        ncrops=n_global + n_local,
        teacher_temp=teacher_temp,
        warmup_teacher_temp=warmup_teacher_temp,
        warmup_teacher_temp_epochs=warmup_teacher_temp_epochs,
        nepochs=n_epochs,
        student_temp=student_temp,
        center_momentum=center_momentum,
    )

    # Create dummy single layer network. We use this to verify
    # that the gradient backprop works properly.
    teacher = torch.nn.Linear(output_dim, output_dim)
    deactivate_requires_grad(teacher)
    student = torch.nn.Linear(output_dim, output_dim)
    orig_teacher = copy.deepcopy(teacher)
    orig_student = copy.deepcopy(student)

    optimizer = torch.optim.SGD(student.parameters(), lr=1)
    orig_optimizer = torch.optim.SGD(orig_student.parameters(), lr=1)

    # Create fake output
    teacher_out = _generate_output(
        batch_size=batch_size,
        n_views=n_global,
        output_dim=output_dim,
        seed=0,
    )
    student_out = _generate_output(
        batch_size=batch_size,
        n_views=n_global + n_local,
        output_dim=output_dim,
        seed=1,
    )

    # Clone input tensors
    orig_teacher_out = torch.cat(teacher_out)
    orig_teacher_out = orig_teacher_out.detach().clone()
    orig_student_out = torch.cat(student_out)
    orig_student_out = orig_student_out.detach().clone()

    # Forward pass
    teacher_out = [teacher(view) for view in teacher_out]
    student_out = [student(view) for view in student_out]
    orig_teacher_out = orig_teacher(orig_teacher_out)
    orig_student_out = orig_student(orig_student_out)

    # Calculate loss
    teacher_temp = orig_loss_fn.teacher_temp_schedule[epoch]
    loss = loss_fn(
        teacher_out=teacher_out,
        student_out=student_out,
        teacher_temp=teacher_temp,
    )
    orig_loss = orig_loss_fn(
        student_output=orig_student_out,
        teacher_output=orig_teacher_out,
        epoch=epoch,
    )

    # Backward pass and optimizer step
    optimizer.zero_grad()
    orig_optimizer.zero_grad()
    loss.backward()
    orig_loss.backward()
    optimizer.step()
    orig_optimizer.step()

    # Loss and loss center should be equal
    center = loss_fn.center.squeeze()
    orig_center = orig_loss_fn.center.squeeze()
    assert torch.allclose(center, orig_center)
    assert torch.allclose(loss, orig_loss)

    # Parameters of network should be equal after backward pass
    for param, orig_param in zip(student.parameters(), orig_student.parameters()):
        torch.testing.assert_close(param, orig_param)
    for param, orig_param in zip(teacher.parameters(), orig_teacher.parameters()):
        torch.testing.assert_close(param, orig_param)
Runs test with the given input parameters.
tests/loss/test_dino_loss.py