Storing tensors for later computation out of memory error


I have a task where at each step I need to encode a Graph Neural Network. Due to the structure of my graph I know that at each step only certain nodes will need re-encoded and in principle I can store the encodings of the other nodes from previous computations for use in this prediction. Once the graph is encoded I run a simple feedforward neural network to output a decision on the Graph Nodes. When I try to store tensors from my past evaluations and update only certain nodes very quickly I run out of memory (64GB) . I feel this problem is somehow due to many computation graphs being stored but I am not sure exactly what is going on. I have also tried to just re-encode from scratch the graph each time. This doesn’t cause any oom issues however its major downside is the time it takes and I want to take advantage of the fact that I already know the encodings for some nodes from past steps.

I am looking for some advice as to what is going on here.