diff --git a/setup.py b/setup.py
index 6baaa4a5..66348335 100644
--- a/setup.py
+++ b/setup.py
@@ -48,6 +48,7 @@
     "timesfm[torch]>=1.2.0,<1.5;python_version>='3.11'",
     "jax;python_version>='3.11'",
 
+    'wrapt>=1.14,<1.15',
 ]