As per https://github.com/logicalclocks/maggy/blob/master/maggy/ablation/ablationstudy.py , we need to write a custom function for the dataset generator to make AblationStudy to work on databricks.