fix: jax gpu note in readme

automl · May 30, 2024 · 4d4a074 · 4d4a074
1 parent 7c4f25e
commit 4d4a074
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -43,7 +43,7 @@ The ARLBench is a benchmark for HPO in RL - evaluate your HPO methods fast and o
 There are currently two different ways to install ARLBench:
 
 <details>
-<summary>After acceptance: PyPI</summary>
+<summary>PyPI</summary>
 You can install ARLBench using `pip`:
 
 ```bash
@@ -85,6 +85,8 @@ make install-envpool
 > [!CAUTION]
 > Windows is currently not supported and also not tested. We recommend using the Linux subsytem if you're on a Windows machine.
 
+If you want to run the ARLBench on GPU, we recommend you check out the [JAX installation guide](https://jax.readthedocs.io/en/latest/installation.html) to see how you can install the correct version for your GPU setup.
+
 ## Quickstart
 
 Here are the two ways you can use ARLBench: via the command line or as an environment. To see them in action, take a look at our [examples](https://github.com/automl/arlbench/tree/main/examples).

diff --git a/pyproject.toml b/pyproject.toml
@@ -8,8 +8,8 @@ dependencies = [
   "codecarbon",
   "memray",
   "pathos",
-  "jax>=0.4.26",
-  "jaxlib>=0.4.26",
+  "jax==0.4.26",
+  "jaxlib==0.4.26",
   "flashbax>=0.1.2",
   "gymnax>=0.0.6",
   "ConfigSpace",