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()
scaler = torch.amp.GradScaler(device_type, enabled=(amp_dtype == torch.float16))