import torch
from torch import nn


torch.manual_seed(7)

device_type = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_type)
amp_dtype = torch.float16 if device_type == "cuda" else torch.bfloat16

if device_type == "cpu":
    torch.backends.mkldnn.enabled = False

if not torch.amp.autocast_mode.is_autocast_available(device_type):
    raise SystemExit(f"autocast is not available for {device_type}")

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 1),
).to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
loss_fn = nn.MSELoss()

inputs = torch.randn(16, 4, device=device)
target = torch.randn(16, 1, device=device)
scaler = torch.amp.GradScaler(device_type, enabled=(device_type == "cuda"))

optimizer.zero_grad(set_to_none=True)
before = model[0].weight.detach().clone()

with torch.amp.autocast(device_type=device_type, dtype=amp_dtype):
    prediction = model(inputs)
    loss = loss_fn(prediction, target)

scaler.scale(loss).backward()
grad_dtype = model[0].weight.grad.dtype
scaler.step(optimizer)
scaler.update()

weight_changed = not torch.equal(before, model[0].weight.detach())

print(f"torch_version={torch.__version__}")
print(f"device_type={device_type}")
print(f"autocast_available={torch.amp.autocast_mode.is_autocast_available(device_type)}")
print(f"amp_dtype={amp_dtype}")
print(f"scaler_enabled={scaler.is_enabled()}")
print(f"prediction_dtype={prediction.dtype}")
print(f"loss_dtype={loss.dtype}")
print(f"grad_dtype={grad_dtype}")
print(f"weight_changed={weight_changed}")
print(f"loss_value={loss.detach().float().item():.6f}")
print("step_complete=True")