Sure here it is. As mentioned my model works when I train a new model from scratch. It is only when continuing my good model that i get the NaNs
from torch.cuda import amp
def mini_trainfp16(model,opt,scheduler,epochs,dataloaders,image_tracker=None,saver=None,saver2=None):
    print('------------------------------------------ Using fp16')
    scaler = amp.GradScaler()
    for epoch in range(epochs):
        for e in dataloaders.keys():
            model.train() if e=='train' else model.eval()
            curr_loader = tqdm(dataloaders[e], total=int(len(dataloaders[e])))
            for idx,data in enumerate(curr_loader):
                x = data['image'].float().to(device);labels_v=data['label'].double().to(device) 
                with amp.autocast():
                    out = model(x)
                    loss = muti_bce_loss_fusionfp16(out, labels_v) 
                
                if e=='train':
                    scaler.scale(loss).backward()
                    scaler.step(opt)
                    scaler.update()
                    opt.zero_grad()
                    scheduler.step()