fridom.framework.utils.jax_utils.free_memory#
- fridom.framework.utils.jax_utils.free_memory()[source]#
This function deletes all live buffers in the JAX backend.
Description#
This function destroys all live buffers in the JAX backend. This is useful for rerunning the code in the same session without running out of memory. Note that the memory is only freed within JAX, not in the operating system. The operating system will still show the same memory usage.