Saving predicted results consumes too much memory

linuxonly801: <p>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.</p>
<pre data-code-wrap="python"><code class="lang-python">from ImmuneBuilder import NanoBodyBuilder2 
predictor = NanoBodyBuilder2() 
for i in range(1000): 
    sequence = {'H': 'QVQLVESGGGLVQPGESLRLSCAASGSIFGIYAVHWFRMAPGKEREFTAGFGSHGSTNYAASVKGRFTMSRDNAKNTTYLQMNSLKPADTAVYYCHALIKNELGFLDYWGPGTQVTVSS'} 
    antibody = predictor.predict(sequence) 
    antibody.save("test.pdb")
</code></pre>
<p><a href="https://private-user-images.githubusercontent.com/69144105/395622041-02371b92-3c34-4514-831e-b7f0a521dc2c.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MzQxNTE4NzAsIm5iZiI6MTczNDE1MTU3MCwicGF0aCI6Ii82OTE0NDEwNS8zOTU2MjIwNDEtMDIzNzFiOTItM2MzNC00NTE0LTgzMWUtYjdmMGE1MjFkYzJjLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDEyMTQlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQxMjE0VDA0NDYxMFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTFlZjQ0OThhNDY4YTlkNjhiYzA4MWJjNDEyNjJkZTMwNDkzM2ZjYTY4MTk4NjQ3ZWViZWI3YmJmNDJhNjczZWYmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0In0.p12LisEq8fpr3dCfXjsDA_vnEeX_POdoex2uLJYuHYs" rel="noopener nofollow ugc"><img src="https://private-user-images.githubusercontent.com/69144105/395622041-02371b92-3c34-4514-831e-b7f0a521dc2c.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MzQxNTE4NzAsIm5iZiI6MTczNDE1MTU3MCwicGF0aCI6Ii82OTE0NDEwNS8zOTU2MjIwNDEtMDIzNzFiOTItM2MzNC00NTE0LTgzMWUtYjdmMGE1MjFkYzJjLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDEyMTQlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQxMjE0VDA0NDYxMFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTFlZjQ0OThhNDY4YTlkNjhiYzA4MWJjNDEyNjJkZTMwNDkzM2ZjYTY4MTk4NjQ3ZWViZWI3YmJmNDJhNjczZWYmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0In0.p12LisEq8fpr3dCfXjsDA_vnEeX_POdoex2uLJYuHYs" alt="2024-12-13_23-59" width="690" height="242"></a></p>
<p>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.</p>
<pre data-code-wrap="python"><code class="lang-python">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
</code></pre>
<p><strong>Desktop (please complete the following information):</strong></p>
<ul>
<li>OS: Ubuntu 20.04</li>
<li>Python version: 3.12</li>
<li>PyTorch version: 2.5</li>
</ul>

ptrblck: <p>Try to <code>detach</code> the output before saving it.</p>