Budget clean task notebook not working

[Jn-Huang]

Thank you so much for developing this tool. I noticed a problem when I try to run the notebook for budget cleaning task (https://github.com/data-centric-ai/dcbench/blob/main/notebooks/dcbench_budgetclean.ipynb). I have the following error:

I am extremely grateful if you can take a look at this problem. Thanks!