Hello all, I am trying to train a model using the automatic mixed precision library by nvidia so here is how the code looks like :
def train(train_dataset, train_dataloader, model, device, optimizer, scheduler):
    model.train()
    train_final_loss = utils.AverageMeter() # This is used to calculate the running averag of loss accross batches
    tk0 = tqdm(train_dataloader, total=len(train_dataloader))

    model.zero_grad()
    for batch_index, data in enumerate(tk0):
        
        images = torch.stack(data[IDX_IMAGES]).to(device, dtype=torch.float32)
        targets = data[IDX_TARGETS]
        image_ids = data[IDX_IMAGE_IDS]

        batch_size = images.shape[0]

        boxes = [target['boxes'].to(device, dtype=torch.float32) for target in targets]
        labels = [target['labels'].to(device, dtype=torch.float32) for target in targets]
        target_res = {'cls' : labels, 'bbox' : boxes}

        outputs = model(images, target_res)
        loss = outputs['loss']
        loss = loss / Config.ACCUMULATE 
        train_final_loss.update(loss.item(), batch_size)
        if Config.USE_GRADIENT_ACCUMULATION:
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            loss.backward()
        
        if (batch_index + 1) % Config.ACCUMULATE == 0 : 
            optimizer.step()
            optimizer.zero_grad()
        tk0.set_postfix(loss=train_final_loss.avg)

    return train_final_loss.avg


I was seeing that the functions need to have its loss divided the accumulation steps. is that reallly necessary? from my understanding, the gradient accumulation will collect the gradients for accumulation steps and then apply it to the network at the end of that number of step, hence it becomes wrong in dividing the loss by ACCUMULATION_STEPS is wrong. please let me know if I am in right direction.