tracmass.free_device_memory

tracmass.free_device_memory()

Release GPU/device memory held by tracmass between runs.

Clears the JAX JIT compilation cache (which pins kmt and dxdy as device constants), runs Python garbage collection to drop any remaining JAX array references, and waits for all asynchronous device operations to complete so memory is actually reclaimed before the next run starts.

Call this in an interactive session (IPython / Jupyter) after time_dependent() returns to avoid GPU memory accumulating across successive runs:

result = tracmass.time_dependent(state, grid, lazy, ...)
tracmass.free_device_memory()   # GPU back to baseline

Notes

This also evicts compiled kernels, so the next call to time_dependent() will retrigger JIT compilation (typically a few seconds). It has no effect on CPU-only runs.

Return type:

None