diff --git a/TrainerClass.py b/TrainerClass.py
index 6babaa5..8518bce 100644
--- a/TrainerClass.py
+++ b/TrainerClass.py
@@ -5,7 +5,8 @@ import matplotlib.pyplot as plt
 import matplotlib
 matplotlib.rcParams['text.usetex'] = True
 
-from sklearn.model_selection import StratifiedGroupKFold, KBinsDiscretizer, StratifiedShuffleSplit, GridSearchCV
+from sklearn.preprocessing import KBinsDiscretizer
+from sklearn.model_selection import StratifiedGroupKFold, StratifiedShuffleSplit, GridSearchCV
 from sklearn.metrics import mean_squared_error, mean_absolute_error
 from sklearn.preprocessing import MinMaxScaler