Saving predicted results consumes too much memory

I want to use NanoBodyBuilder2 to predict a lot of structures (about 1000), but when I use a loop to iterate my sequences, something like the following code, the GPU memory increased every time i got a structure. And finally run OOM.

from ImmuneBuilder import NanoBodyBuilder2 
predictor = NanoBodyBuilder2() 
for i in range(1000): 
    sequence = {'H': 'QVQLVESGGGLVQPGESLRLSCAASGSIFGIYAVHWFRMAPGKEREFTAGFGSHGSTNYAASVKGRFTMSRDNAKNTTYLQMNSLKPADTAVYYCHALIKNELGFLDYWGPGTQVTVSS'} 
    antibody = predictor.predict(sequence) 
    antibody.save("test.pdb")

2024-12-13_23-59

But, if i modify the code like below, the GPU memory stays the same for predicting the result. But when running ab.save(f"test{cnt}.pdb"), the GPU memory increases very fast, and finally got OOM.

from ImmuneBuilder import NanoBodyBuilder2 
predictor = NanoBodyBuilder2() 
results = [] 
for i in range(1000): 
    sequence = {'H': 'QVQLVESGGGLVQPGESLRLSCAASGSIFGIYAVHWFRMAPGKEREFTAGFGSHGSTNYAASVKGRFTMSRDNAKNTTYLQMNSLKPADTAVYYCHALIKNELGFLDYWGPGTQVTVSS'} 
    antibody = predictor.predict(sequence) 
    results.append(antibody) 

cnt = 0 
for ab in results:
    ab.save(f"test{cnt}.pdb") 
    cnt += 1

Desktop (please complete the following information):

  • OS: Ubuntu 20.04
  • Python version: 3.12
  • PyTorch version: 2.5

Try to detach the output before saving it.