Inference exhausts GPU memory

[wkguo]

As we deploy fine-tuned Pi0 on a local commercial GPU with 48 GB memory,
it seems to take around 37 GB of memory to successfully load the model into the server, when running the uv run openpi/scripts/serve_policy.py ...
As the repo reports a recommendation of >8GB memory, may I ask if there are any other settings to reduce the real-world deployment cost?

[kvablack]

JAX preallocates 75% of GPU memory by default for more efficient allocation. See this page:

https://docs.jax.dev/en/latest/gpu_memory_allocation.html