Wouldn’t this work?
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Creates a GradScaler once at the beginning of training.
scaler = GradScaler(enabled=use_amp)

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        
        with autocast(enabled=use_amp):
            output = model(input)
            loss = loss_fn(output, target)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()