Update train_model.py

src/models/train_model.py

@@ -29,6 +29,9 @@ def train_model():
         train_df=train_df,
         eval_df=eval_df,
         batch_size=params["batch_size"],
+        max_epochs=params["epochs"],
+        use_gpu=params["use_gpu"],
+        learning_rate=float(params["learning_rate"]),
 
     model.save_model(model_dir=params['model_dir'])