revert pred index to tree/forest_ite_pred

uber · Nov 28, 2023 · 9040882 · 9040882
1 parent c86d0c6
commit 9040882
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_causal_trees.py b/tests/test_causal_trees.py
@@ -84,7 +84,7 @@ def test_fit(self, generate_regression_data):
             treatment_col="is_treated",
             treatment_effect_col="treatment_effect",
         )
-        assert df_qini["pred"] > df_qini["Random"]
+        assert df_qini["ctree_ite_pred"] > df_qini["Random"]
 
     @pytest.mark.parametrize("return_ci", (False, True))
     @pytest.mark.parametrize("bootstrap_size", (500, 800))
@@ -165,7 +165,7 @@ def test_fit(self, generate_regression_data, n_estimators):
             treatment_col="is_treated",
             treatment_effect_col="treatment_effect",
         )
-        assert df_qini["pred"] > df_qini["Random"]
+        assert df_qini["crforest_ite_pred"] > df_qini["Random"]
 
     @pytest.mark.parametrize("n_estimators", (5,))
     def test_predict(self, generate_regression_data, n_estimators):