Inference does not use GPUs.

I have updated Jax to the version support CUDA.
```
pip install --upgrade --user jax==0.4.6 jaxlib==0.4.6+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
```

I also check that Jax now can detect GPU.

```
(alphageometry) [jpark3234@atl1-1-03-010-10-0 alphageometry]$ python -c "import jax; print(jax.devices())"
[StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0)]
```

But when I run the code in run.sh, it does not utilize GPU (only few memcopy occurred in GPU).

Anyone has met/solved this problem?