fix jax installation

epignatelli · Mar 29, 2021 · 919f206 · 919f206
1 parent 2dfa6cb
commit 919f206
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/install_jax.sh b/install_jax.sh
@@ -1,2 +1,2 @@
-CUDA_VERSION=$(nvidia-smi | grep -o 'CUDA Version: [0-9].\.[0-9]' | sed 's/.*: //')
+CUDA_VERSION=$(nvidia-smi | grep -o 'CUDA Version: [0-9].\.[0-9]' | sed 's/.*: //' | sed 's/\.//')
 pip install --upgrade jax jaxlib==0.1.64+cuda$CUDA_VERSION -f https://storage.googleapis.com/jax-releases/jax_releases.html/