How to set specific GPU?

[MichaelMMeskhi]

I have 2 GPUs but want Jax to utilize only the second one. Only solution I have found to work is:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1" # second gpu

import jax
import jax.numpy as jnp
...

Is there a better way to do this?

Jax sees the following:

>>> jax.devices()
[StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=1, process_index=0, slice_index=0)]

[selamw1]

You can directly specify the desired devices within your JAX code using the jax.devices function. This approach offers more control and clarity:

>>> import jax
>>> 
>>> # Retrieve available devices
>>> available_devices = jax.devices()
>>> print(f"Available devices: {available_devices}")
Available devices: [cuda(id=0), cuda(id=1)]
>>> 
>>> # Select the second GPU
>>> second_gpu = available_devices[1]
>>> print(f"Using device: {second_gpu}")
Using device: cuda:1
>>> 
>>> # Define a function to run on the second GPU
>>> def square_on_gpu(x):
...   return x * x
... 
>>> # Compile the function for the second GPU
>>> square_gpu = jax.jit(square_on_gpu, device=second_gpu)
>>> 
>>> # Perform a computation on the second GPU
>>> result = square_gpu(4.0)
>>> print(f"Result: {result}")  # Output: Result: 16.0
Result: 16.0