After following @ptrblck’s suggestions here’s how my train() function is looking like (consider this to be the launcher expected by torch.multiprocessing.spawn().
def train(rank, num_epochs, world_size):
    init_process(rank, world_size)
    torch.manual_seed(0)
    
    model = create_model()
    torch.cuda.set_device(rank)
    model.cuda(rank)
    model = DistributedDataParallel(model, device_ids=[rank])
    swa_model = torch.optim.swa_utils.AveragedModel(model)

    learning_rate = 1e-3
    optimizer = torch.optim.Adam(model.parameters(), 
                             lr=learning_rate * world_size)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
    swa_start = 10
    swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, swa_lr=0.05)
    criteria = nn.CrossEntropyLoss()
    scaler = torch.cuda.amp.GradScaler(enabled=True)

    
    train_loader, val_loader = get_dataloader(rank, world_size)
    
    for epoch in range(num_epochs):
    	model.train()
        for batch in train_loader:
            with torch.cuda.amp.autocast(enabled=True):
                image = batch['image'].cuda(rank, non_blocking=True)
                mask = batch['mask'].cuda(rank, non_blocking=True)
                
                pred = model(image)
                
                loss = criteria(pred, mask.unsqueeze(1))

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

Things are now working as expected. @ptrblck if I could do anything better to further optimize the performance please let me know.
Majority of the SWA code comes from the official docs 1.