diff --git a/src/trainScript.py b/src/trainScript.py
index b873529..19cf262 100644
--- a/src/trainScript.py
+++ b/src/trainScript.py
@@ -59,7 +59,7 @@ linear_model.compile(loss = tf.losses.MeanSquaredError(),
                       optimizer = tf.optimizers.Adam(1))
 
 # train model
-history = linear_model.fit(house_price_features, house_price_labels, epochs=numberOfEpoch, )
+history = linear_model.fit(house_price_features, house_price_labels, epochs=int(numberOfEpoch), )
 #callbacks=[cp_callback])
 # print(history)