diff --git a/README.md b/README.md
index 070f00d019..8ca186dcfc 100644
--- a/README.md
+++ b/README.md
@@ -75,8 +75,13 @@ macOS, run:
 pip install mlx
 ```
 
-On Linux, the same command (``pip install mlx``) will install the CUDA backend
-by default. To install a CPU-only Linux package, run:
+To install the CUDA backend on Linux, run:
+
+```bash
+pip install mlx[cuda]
+```
+
+To install a CPU-only Linux package, run:
 
 ```bash
 pip install mlx[cpu]
diff --git a/docs/src/install.rst b/docs/src/install.rst
index 7d9f6d90de..1e7a015ca4 100644
--- a/docs/src/install.rst
+++ b/docs/src/install.rst
@@ -30,7 +30,7 @@ MLX has a CUDA backend which you can install with:
 
 .. code-block:: shell
 
-    pip install mlx
+    pip install mlx[cuda]
 
 To install the CUDA package from PyPi your system must meet the following
 requirements:
diff --git a/setup.py b/setup.py
index 4492a20c2b..d15f36b2b7 100644
--- a/setup.py
+++ b/setup.py
@@ -274,13 +274,11 @@ if __name__ == "__main__":
     #  - Package name is back-end specific, e.g mlx-metal
     if build_stage != 2:
         if build_stage == 1:
-            install_requires += [
-                f'mlx-metal=={version}; platform_system == "Darwin"',
-                f'mlx-cuda=={version}; extra != "cpu" and platform_system == "Linux"',
-            ]
-            extras["cpu"] = [
-                f'mlx-cpu=={version}; extra == "cpu" and platform_system == "Linux"'
-            ]
+            install_requires.append(
+                f'mlx-metal=={version}; platform_system == "Darwin"'
+            )
+            extras["cuda"] = [f'mlx-cuda=={version}; platform_system == "Linux"']
+            extras["cpu"] = [f'mlx-cpu=={version}; platform_system == "Linux"']
 
         _setup(
             name="mlx",