Hi, thanks for your reply. I am also using AMP (also using SWA). Here’s the skeleton of what I’m doing:
model = create_model(args)
model = nn.DataParallel(model)
swa_model = torch.optim.swa_utils.AveragedModel(model)
scaler = torch.cuda.amp.GradScaler()
optimizer = torch.optim.AdamW(model.parameters(), lr=7e-3)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer,
                                max_lr=7e-2,
                                epochs=5, steps_per_epoch=len(dataloader), 
                                pct_start=0.3)  
swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, 
                                            swa_lr=0.5, 
                                            anneal_epochs=5)
num_epochs = 20
swa_start = 15

for epoch in (1, range(num_epochs)+1):
    for imgs,labels in dataloader:

        with torch.cuda.amp.autocast():
            output = model(imgs)
            loss = loss_fn(output, labels)

        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        
        if epoch < swa_start:
            scheduler.step()
            
    if epoch >= swa_start:
        swa_model.update_parameters(model)
        swa_scheduler.step()