how to set the maximum gpu memory occupied by jax project on a single GPU,

[andaolong]

2023-9-18_11:03:20
I want to know how to set the maximum gpu memory occupied by jax project on a single GPU, so that I can use the GPU of our lab public server conveniently.

[anh-tong]

You can refer to the document here GPU memory allocation and set environment variable as

export XLA_PYTHON_CLIENT_MEM_FRACTION=.XX

where XX is the percentage of GPU you want to use.

[andaolong]

Thank you very much for your advice. Your advice helped me solve the problem. Thank you again for your help.