This is the step up I use -
scaler = GradScaler()
loss = DiceFocalLoss()
for _, data in enumerate(dataloader):
    break

image = torch.randn((1, 3, 640, 640)).cuda()
mask = torch.randint(0, 2, (1, 640, 640)).cuda().float()

model = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=1, padding=1),
                       nn.Conv2d(in_channels=3, out_channels=1, kernel_size=1, stride=1, padding=0))
model = model.cuda().train()

optim = torch.optim.Adam(model.parameters(), lr=1e-3)
use_amp = False
loss_list = []

for i in range(10):
    #optim.zero_grad()
    with torch.cuda.amp.autocast(enabled=use_amp):
        out = model(image)
        l = loss(out, mask)
    if use_amp:
        scaler.scale(l).backward()
        scaler.step(optim)
        scaler.update()
    else:
        l.backward()
        optim.step()

    print(l)

print("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-==-==-=-==-==-=-==-==-=-=-==-")
model1  = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=1, padding=1),
                       nn.Conv2d(in_channels=3, out_channels=1, kernel_size=1, stride=1, padding=0))

use_amp = True
model1.cuda().train()
for i in range(10):
    #optim.zero_grad()
    with torch.cuda.amp.autocast(enabled=use_amp):
        out = model1(image)
        l = loss(out, mask)
    if use_amp:
        scaler.scale(l).backward()
        scaler.step(optim)
        scaler.update()
    else:
        l.backward()
        optim.step()

    print(l)


The output is -
tensor(2.2960, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.2791, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.2631, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.2475, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.2321, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.2169, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.2018, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.1868, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.1719, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.1571, device='cuda:0', grad_fn=<AddBackward0>)
=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-==-==-=-==-==-=-==-==-=-=-==-
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)
tensor(2.8875, device='cuda:0', grad_fn=<AddBackward0>)

Process finished with exit code 0


The top one is FP32 and the below one FP16.
There is no change in losses in FP16. One should expect the same behavior right? like losses might be different but it should decrease.  The loss is Dice Focal loss as described above.