Thanks for the update.



 jabhinav:

Solutions tried: I have tried gradient accumulation where I am splitting batch_size into sub-batches, calling loss.backward for each of them before calling optimiser.step().


If your model is not depending on the batch size, e.g. via batchnorm layers, this should be a valid approach to trade compute for memory.
Did you still run out of memory using this approach?