Unable to use gpu: might be a jax(or jaxlib) issue

[utkarshp1161]

When I am running the gpax_vidkl_eels.ipynb I am getting below warning:

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to CPU.

How to reproduce:

I created the python env using following steps: conda create -n gpax_hae python==3.10 conda activate gpax_hae pip install -e . (inside gpax repository) ------- this install most of the dependencies including jax and jaxlib pip install jupyter pip install ipykernel

[ziatdinovmax]

@utkarshp1161 please try installing following these instructions: https://github.com/google/jax#instructions

[utkarshp1161]

It works now. Thanks.