Changing fit to fit_transform