Thanks. I cannot reproduce the issue using this code:
# your model defnitions

print(torch.cuda.get_device_name(0))
torch.backends.cudnn.benchmark = True

model = UNet(2).cuda()
x = torch.randn(24, 2, 256, 256, device='cuda')
target = torch.randn(24, 2, 256, 256, device='cuda')
out = model(x)
print(out.shape)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999))

criterion = nn.MSELoss()

scaler = torch.cuda.amp.GradScaler()

print('running amp')
for _ in range(10):
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        out = model(x)
        loss = torch.mean(torch.abs(out - target))
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    print(out.shape)

using an RTX2080Ti with PyTorch 1.6.0+CUDA10.2, 1.7.1+CUDA10.2, 1.7.1+CUDA11.0.
Are you seeing the crash using a single GPU as well or only in this particular setup with two different GPUs?