import time
import torch
from torch.optim import Adam, SGD, Adagrad, NAdam

optim_cls = Adam
NUM_PARAMS = 200
kwargs = { "lr": 0.01, "foreach": True }

torch._dynamo.reset()
# torch._inductor.metrics.reset()
input = torch.ones([10, 10], device="cuda:0")
model = torch.nn.Sequential(
    *[torch.nn.Linear(10, 10, device="cuda:0") for _ in range(NUM_PARAMS)]
)

model(input).sum().backward()
opt_compiled = optim_cls(model.parameters(), **kwargs)
compiled_step = torch.compile(opt_compiled.step)

with torch.set_grad_enabled(False):
    start_time = time.time()
    compiled_step()
    print("compile opt took: %s seconds", time.time() - start_time)

print(torch._dynamo.utils.compile_times())

# Handle complex parameters
for i in range(len(device_grads)):
    if torch.is_complex(device_params[i]):
        device_grads[i] = torch.view_as_real(device_grads[i])
        device_exp_avgs[i] = torch.view_as_real(device_exp_avgs[i])
        device_exp_avg_sqs[i] = torch.view_as_real(device_exp_avg_sqs[i])
        device_max_exp_avg_sqs[i] = torch.view_as_real(device_max_exp_avg_sqs[i])
        device_params[i] = torch.view_as_real(device_params[i])