trax.math.device_count() returning 1 whereas I have 4 GPU Tesla V100 about trax HOT 2 CLOSED

Problem solved. It was problem with jax. I reinstall using the following instruction

# install jaxlib
PYTHON_VERSION=cp37  # alternatives: cp27, cp35, cp36, cp37
CUDA_VERSION=cuda101  # alternatives: cuda90, cuda92, cuda100, cuda101
PLATFORM=linux_x86_64  # alternatives: linux_x86_64
BASE_URL='https://storage.googleapis.com/jax-releases'
pip install --upgrade $BASE_URL/$CUDA_VERSION/jaxlib-0.1.37-$PYTHON_VERSION-none-$PLATFORM.whl

pip install --upgrade jax  # install jax(base)

from trax.

hi @sanjibnarzary, have you ever encountered this error:
Can't find ptxas binary in ${CUDA_DIR}/bin. Will back to the GPU driver for PTX -> sass compilation. This is OK so long as you don't see a warning below about an out-of-date driver version.

The ptxas is in the way /usr/local/cuda/bin, but trax doesn't seems to find it.

I tried your instructions, but it didn't work.

from trax.