From fda28a35b60a42e98f818878b35f8e7b0f41e54d Mon Sep 17 00:00:00 2001
From: WLS2002 <64534280+WLS2002@users.noreply.github.com>
Date: Sat, 13 Jul 2024 21:28:53 +0800
Subject: [PATCH] Update README.md

---
 README.md | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index a285ec4..9b8d0ea 100644
--- a/README.md
+++ b/README.md
@@ -37,7 +37,22 @@ TensorNEAT is a JAX-based libaray for NeuroEvolution of Augmenting Topologies (N
     - Test neuroevolution algorithms on advanced **RL tasks** (Brax, Gymnax).
 
 ## Installation
-Install `tensorneat` from the GitHub source code:
+
+1. Install the correct version of [JAX](https://github.com/google/jax). We recommend `jax >= 0.4.28`.
+
+For cpu version only, you may use:
+```
+pip install -U jax
+```
+
+For nvidia gpus, you may use:
+```
+pip install -U "jax[cuda12]"
+```
+For details of installing jax, please check https://github.com/google/jax.
+
+
+2. Install `tensorneat` from the GitHub source code:
 ```
 pip install git+https://github.com/EMI-Group/tensorneat.git
 ```