Fix JAX GPU tests.

Upgraded JAX version to 0.6.2.

Author: hertschuh

requirements-jax-cuda.txt

@@ -8,7 +8,7 @@ torch==2.6.0
 
 # Jax with cuda support.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-jax[cuda12]==0.6.0
+jax[cuda12]==0.6.2
 flax
 
 -r requirements-common.txt