diff --git a/econml/tests/test_shap.py b/econml/tests/test_shap.py
index d2c6153f9..2c69b39d4 100644
--- a/econml/tests/test_shap.py
+++ b/econml/tests/test_shap.py
@@ -181,7 +181,9 @@ def eta_sample(n):
 
         # test shap could generate the plot from the shap_values
         heatmap(shap_values1["Y0"]["orange"], show=False)
-        waterfall(shap_values1["Y0"]["orange"][6], show=False)
+        if shap.__version__ != "0.40.0":
+            # waterfall is broken in this version, fixed by https://github.com/slundberg/shap/pull/2444
+            waterfall(shap_values1["Y0"]["orange"][6], show=False)
         scatter(shap_values1["Y0"]["orange"][:, "A"], show=False)
         bar(shap_values1["Y0"]["orange"], show=False)
         beeswarm(shap_values1["Y0"]["orange"], show=False)
diff --git a/setup.cfg b/setup.cfg
index c816f9d65..7761ff78c 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -39,7 +39,7 @@ install_requires =
     joblib >= 0.13.0
     statsmodels >= 0.10
     pandas
-    shap >= 0.38.1, < 0.40.0
+    shap >= 0.38.1, < 0.41.0
     dowhy < 0.8
     lightgbm
 test_suite = econml.tests