diff --git a/TrainerClass.py b/TrainerClass.py
index cb351bf..0529a5a 100644
--- a/TrainerClass.py
+++ b/TrainerClass.py
@@ -15,6 +15,7 @@ from xgboost import XGBRegressor
 
 import keras
 from keras import layers
+import ray
 from ray import tune
 from ray.tune.schedulers import ASHAScheduler