More prominent JAX GPU instructions

README.md

@@ -54,6 +54,9 @@ conda create -n arlbench python=3.10
 conda activate arlbench
 ```
 
+The instructions below will help you install the default version of ARLBench with the CPU version of JAX. 
+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 before proceeding.
+
 <details>
 <summary>PyPI</summary>
 You can install ARLBench using `pip`:
@@ -92,8 +95,6 @@ make install-envpool
 > [!CAUTION]
 > Windows is currently not supported and also not tested. We recommend using the [Linux subsytem](https://en.wikipedia.org/wiki/Windows_Subsystem_for_Linux) 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).