Device Memory Profiling only explains half of memory usage

[JGameCreation]

HI, I'm trying to use Device Memory Profiling (https://jax.readthedocs.io/en/latest/device_memory_profiling.html) to analyze my memory usage. However, the pprof graph only shows a total VRAM usage of 4700MB while nvidia-smi reports 10800. Where could this difference be coming from?

[YouJiacheng]

JAX will preallocate 90% of currently-available GPU memory when the first JAX operation is run.

https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
Did you set XLA_PYTHON_CLIENT_ALLOCATOR=platform when you used nvidia-smi?