demo code running uses up system RAM in colab

Running the code on google colab under A100 GPU with 80G system RAM. The running does not take too much GPU RAM but the system RAM was taken by first 40G and then 80G before crash.

<img width="1701" alt="Image" src="https://github.com/user-attachments/assets/aab90a39-db71-4bc2-bb69-2a4d98a881be" />