diff --git a/ium_predict.py b/ium_predict.py
index a2617cb..85c564d 100644
--- a/ium_predict.py
+++ b/ium_predict.py
@@ -26,7 +26,7 @@ def get_x_y(data):
 
 
 def predict():
-    model = load_model('baltimore_model3')
+    model = load_model('baltimore_model.h5')
 
     train = pd.read_csv('baltimore_train.csv')
     baltimore_data_test = pd.read_csv('baltimore_test.csv')
diff --git a/ium_train.py b/ium_train.py
index 1fb49da..a88966d 100644
--- a/ium_train.py
+++ b/ium_train.py
@@ -48,7 +48,7 @@ def train_model():
         validation_split=0.2)
     hist = pd.DataFrame(history.history)
     hist['epoch'] = history.epoch
-    model.save('baltimore_model3')
+    model.save('baltimore_model.h5')
 
 
 train_model()