Specifying device on a multi GPU machine

[hegde95]

Is it possible to have environments across multiple gpus? or is it even possible to specify which gpu will be used by brax?

[gijskoning]

Brax uses Jax to compile the training loop with XLA. This means it can run on the gpu including the environments. Your question is mostly related to the Jax api. With Jax you can specify the devices to be used: cpu, gpu or multiple GPUs. https://jax.readthedocs.io/en/latest/notebooks/quickstart.html

[hegde95]

For posterity, I had to set the environment variable to make BRAX (and JAX) use a specific GPU