def test_conv_train_loop():
    """
    Assert that training the interval conv layer gives the same results as a regular layer
    """
    output_channels = 12
    input_channels = 3
    loss_fn = torch.nn.MSELoss()
    synth_model = get_synthetic_model()

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    target = torch.rand(size=(32, 12, 21, 21)).to(device)

    synthetic_data = torch.rand(32, input_channels, 25, 25).to(device)
    model = synth_model(
        input_shape=synthetic_data.shape,
        output_channels=output_channels,
        kernel_size=5,
        bias=True,
        stride=1,
        to_debug=False,
    ).to(device)

    class TestModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv2d(
                in_channels=3,
                out_channels=12,
                kernel_size=5,
                bias=True,
                stride=1,
            )

        def forward(self, x):
            return self.conv(x)

    test_model = TestModel()
    # Get the weights we will transfer over
    # Set the weights in the normal model
    test_model.conv.weight = torch.nn.Parameter(
        torch.reshape(
            torch.tensor(model.conv1.conv.weight.data.cpu().detach().numpy()),
            shape=(output_channels, input_channels, 5, 5),
        )
    )
    test_model.conv.bias = torch.nn.Parameter(torch.tensor(model.conv1.bias_to_grad.data.cpu().detach().numpy()))
    test_model.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
    test_opt = torch.optim.Adam(test_model.parameters(), lr=0.0001)

    for _ in range(5):
        model.zero_grad()

        # Test model is for ground truth comparison
        test_model.zero_grad()
        output_from_equivalent = model.forward(synthetic_data)
        loss = loss_fn(output_from_equivalent, target)
        loss.backward()

        equivalent_grads = model.conv1.conv.weight.grad
        equivalent_grads = (
            torch.reshape(equivalent_grads, shape=(output_channels, input_channels, 5, 5)).detach().clone().to(device)
        )
        equivalent_bias = model.conv1.bias_to_grad.grad.clone().to(device)

        output_from_conv = test_model.forward(synthetic_data)
        loss = loss_fn(output_from_conv, target)
        loss.backward()

        assert torch.allclose(equivalent_grads, test_model.conv.weight.grad, atol=1e-05)
        assert torch.allclose(equivalent_bias, test_model.conv.bias.grad, atol=1e-05)

        optimizer.step()
        test_opt.step()

        reshaped_weights = torch.reshape(
            model.conv1.conv.weight.data.clone().detach(), shape=(output_channels, input_channels, 5, 5)
        ).to(device)

        assert torch.allclose(reshaped_weights, test_model.conv.weight.data, atol=1e-05)

        # function to re-pop the dense layer
        model.conv1.re_convert(device)

        # Sanity check! Are the grads still present and the same after the conversion?
        equivalent_grads = model.conv1.conv.weight.grad
        equivalent_grads = (
            torch.reshape(equivalent_grads, shape=(output_channels, input_channels, 5, 5)).detach().clone().to(device)
        )
        equivalent_bias = model.conv1.bias_to_grad.grad

        assert torch.allclose(equivalent_grads, test_model.conv.weight.grad, atol=1e-05)
        assert torch.allclose(equivalent_bias, test_model.conv.bias.grad, atol=1e-05)

    Assert that training the interval conv layer gives the same results as a regular layer
    
tests/estimators/certification/test_interval.py