Hi
Code : https://gist.github.com/amogh112/deecd6d388c43cc58c33ed9c659b4b4d 2
I observed that during training, things were fine until 5th epoch when the CPU usage suddenly shot up(see image for RAM usage).

I made sure that loss was detached before logging.
I observed that this does not happen if I comment out the loss.backward call(Line 181 1)
Finally I was able to get around by collected garbage using gc.collect() after every 50 batches. 5
But it still slows the epoch by a lot.
Some explanation for the code:
Link to model code: I added some linear layers to a base model(VLBert) that was training without any issues. I do a forward pass for the base model, take an embedding from the output and pass it through the linear layers to produce 1 output.
Then BCEWithLogits is then used to calculate the loss wrt labels for this output.
Then backward is called.
The GPU(RTX 2080 TI) usage is not a problem, just the CPU usage. Printing out the tensors using gc.get_objects() as explained here shows there are no tensors on the CPU after every epoch.
What is the reason for garbage not being collected, any issues with the code?