diff --git a/requirements.txt b/requirements.txt
index 5eca6f4..2c1c63f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -25,3 +25,5 @@ sympy==1.12
 torch==2.2.1
 triton==2.2.0
 typing_extensions==4.10.0
+ml-collections==0.1.1
+dm-tree==0.1.8