Modified overfitting.py to make it easier to run on varied data sets