diff --git a/python/pyproject.toml b/python/pyproject.toml
index c757e1297..7c6261983 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -49,6 +49,8 @@ srt = [
     "torch==2.5.1",
     "cuda-python",
     "outlines>=0.0.44,<=0.1.11",
+    "partial_json_parser",
+    "einops",
 ]
 
 # HIP (Heterogeneous-computing Interface for Portability) for AMD