Thanks for your response Alban.
I do not do anything specific at this time. I am using 8 GPUs with apex which uses multiprocessing I believe:

from apex.parallel import DistributedDataParallel as DDP
model = DDP(model, delay_allreduce=True)
optim = AdamW(model.parameters(), lr=args.learning_rate, correct_bias=False, eps=args.adam_epsilon)
model, optim = amp.initialize(model, optim, opt_level=args.opt_level, loss_scale=args.loss_scale,
keep_batchnorm_fp32=args.keep_batchnorm_fp32, verbosity=0)

Since it takes a lot more time to run 5 epochs without using the 8 GPUs, I haven’t run it that long to see if its resolved without using DDP.
Not sure if relevant and a source of reference cycles, I take a pre-trained base model and load its saved weights, and then initialise the weights for added layers.

Did you managed to get any insight of which objects get collected when you run this? Could it be some of the logging logic that creates these reference cycles?

In the code I shared, I do not see any source of reference cycles.  I am simply logging losses with Tensorboard.
SInce removing .backward() does not lead to a leakage, I assumed that there are a few tensors on CPU that are in the computation graph. However upon printing out the tensors on all devices at different stages using this function 2, I found that there are some CPU tensors while batch is being processed in training (as received from dataloader before being moved to cuda in new variables eg - imgs = data["imgs"].cuda()). But printing out the tensors on CPU before the next training epoch, I did not find anything on CPU. Also number of tensors on CPU at start of each batch is the same.