1) Make tqdm import allow error and use it for the grid search as well
2) Allow functionality to change the flatten method in the 2D plot