How to use specific GPU with ott-jax

[prabhant]

Hi,

How can i use GPU_ID example: 7 in ott-jax.
Using it directly via jax.device does not work. CUDA_VISIBLE_DEVICES also does not work. OTT-JAX ends up using only gpuid 0 every time.

[michalk8]

Hi @prabhant ,
setting, e.g., CUDA_VISIBLE_DEVICES='1' before start of the program should be working as expected. What jax/ott-jax versions are you using?