disable_gpu_memory_preallocation

class brainpy.math.disable_gpu_memory_preallocation(release_memory=True)

Disable pre-allocating the GPU memory.

This disables the preallocation behavior. JAX will instead allocate GPU memory as needed, potentially decreasing the overall memory usage. However, this behavior is more prone to GPU memory fragmentation, meaning a JAX program that uses most of the available GPU memory may OOM with preallocation disabled.

Parameters:

release_memory (bool) – bool. Whether we release memory during the computation.