diff --git a/tools/python/requirements.txt b/tools/python/requirements.txt
index db0ec26..09121c8 100644
--- a/tools/python/requirements.txt
+++ b/tools/python/requirements.txt
@@ -26,3 +26,5 @@
 tabulate
 
 casadi>=3.6.6
+
+jax[cuda12]