Hello,
Because of memory constraints, I can only use batch_size of 1. But then I came across a trick called “Gradient Accumulation”. I have implemented two versions of it and would like to which one is correct and why, if possible.
Version 1:
# Backpropogation step
            total_loss.backward()

            # Accumulate Gradients and Update/Back-propagate at every 8th iteration.
            if (iter_i+1) % 32 == 0:
                optimizer.step()
                optimizer.zero_grad()

Version 2:
            # Accumulate Gradients and Update/Back-propagate at every 8th iteration.
            if (iter_i+1) % 32 == 0:
                total_loss.backward()
                optimizer.step()
                optimizer.zero_grad()

The total_loss consists of the following loss functions:
cls_loss_function = HeatmapLoss(reduction='mean')  # Custom Loss
txty_loss_function = nn.BCEWithLogitsLoss(reduction='none')
twth_loss_function = nn.SmoothL1Loss(reduction='none')

Should the total_loss be divided by 32 before using the backward function? I have several BatchNorm2d layers in the architecture.
So the question here is whether the total_loss.backward() should be inside or outside the if statement.