Update README.md

README.md

@@ -58,15 +58,15 @@ We've modified Flax's 'lm1b' example to train on Japanese dataset. You can find
 
 ```
 pip install transformers>=4.39.0
-pip install jax==0.4.13
-pip install flax==0.6.11
+pip install jax==0.4.31
+pip install flax==0.8.3
 pip install sentencepiece==0.1.99
 
 # For CPU
-pip install -U "jax[cpu]==0.4.13"
+pip install -U "jax[cpu]==0.4.31"
 
 # For GPU
-pip install -U "jax[cuda12_pip]==0.4.13" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install -U "jax[cuda12]==0.4.31"
 ```
 
 Note: Set **trust_remote_code=True** to load our custom model.