import os

import torch


def run_once():
    torch.manual_seed(1234)

    model = torch.nn.Sequential(
        torch.nn.Linear(4, 3),
        torch.nn.ReLU(),
        torch.nn.Linear(3, 1),
    )
    optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

    inputs = torch.randn(5, 4)
    targets = torch.randn(5, 1)

    optimizer.zero_grad(set_to_none=True)
    loss = torch.nn.functional.mse_loss(model(inputs), targets)
    loss.backward()
    optimizer.step()

    params = torch.cat([parameter.detach().flatten() for parameter in model.parameters()])
    return round(loss.item(), 8), params


torch.use_deterministic_algorithms(True)
torch.backends.cudnn.benchmark = False

first_loss, first_params = run_once()
second_loss, second_params = run_once()

print(f"deterministic_enabled={torch.are_deterministic_algorithms_enabled()}")
print(f"cudnn_benchmark={torch.backends.cudnn.benchmark}")
print(f"cublas_workspace_config={os.environ.get('CUBLAS_WORKSPACE_CONFIG', 'not set')}")
print(f"loss_1={first_loss:.8f}")
print(f"loss_2={second_loss:.8f}")
print(f"losses_match={first_loss == second_loss}")
print(f"parameters_match={torch.equal(first_params, second_params)}")