Possible Memory Leak or Excessive RAM Usage Error in batch model of colabfold_batch Locally

Dear Developers,

I am experiencing an issue with the colabfold_batch tool when processing a batch of protein sequences LOCALLY, each around 600 amino acids long. Despite having a system (256GB RAM and 24GB GPU memory), the tool appears to be continuously consuming memory and eventually fails due to what seems like a memory-related error.

Current Behavior

When executing colabfold_batch for a batch of protein sequences, the process starts but gradually accumulates memory usage until it crashes. It seems that the application is not efficiently releasing memory between individual runs or there might be a slow-compilation problem causing excessive memory consumption.
I've noticed that the RAM (256GB RAM for me) usage continuously increases over time, even though the GPU memory appears stable. The sequences being processed are relatively moderate in size, which leads me to believe that it's the RAM that cause crash.

Error Log:
The error messages received during execution point towards slow compilation processes:

E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65] 
********************************
[Compiling module jit_apply_fn] Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************
E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 2m23.006026321s

********************************
[Compiling module jit_apply_fn] Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************

These warnings suggest that the compiler is taking an unexpectedly long time, which may contribute to the excessive RAM consumption.

System Configuration & Observations:

• My workstation has 256GB of RAM, ensuring ample capacity for handling multiple instances concurrently.
• The GPU has 24GB of memory, which should be sufficient for typical protein folding tasks (600AA this time).
• During the runs, it is clear that the growth in memory usage is attributed to the system (CPU) RAM rather than the GPU memory.

Request for Assistance:
Is there any configuration setting or optimization method that can be applied to colabfold_batch to allow it to run indefinitely without exhausting system memory? Given the size of my hardware resources, I'm looking for guidance on how to prevent memory leaks or optimize the memory allocation strategy for these medium-sized predictions.

Thank you for considering this issue and providing support.

Best regards
lshi

Start at the colabfold_batch entry point and reproduce the reported batch run while monitoring CPU RAM and GPU memory. Use the provided XLA_FLAGS dump setting to investigate the slow compilation messages; the issue names no source files or tests, and done would require identifying the memory-growth cause or a confirmed configuration-based mitigation.