diff --git a/prediction.py b/prediction.py
index 1cab3a5..a8e19bc 100644
--- a/prediction.py
+++ b/prediction.py
@@ -77,15 +77,15 @@ def get_prediction(smg, sml, indexes, bls, blhs, dis, dps, dphs, model_name):
     w = [0.86, 0.8, 0.06]
     for i in range(1, 4):
         try:
-            model = tf.train.import_meta_graph(
+            model = tf.compat.v1.train.import_meta_graph(
                 model_name + '_v' + str(i) + '/model.ckpt.meta')
         except BaseException:
             print("Something wrong with the model.")
             continue
-        with tf.Session() as sess:
+        with tf.compat.v1.Session() as sess:
             try:
                 model.restore(sess, model_name + '_v' + str(i) + "/model.ckpt")
-                graph = tf.get_default_graph()
+                graph = tf.compat.v1.get_default_graph()
                 synmgt, synmlt, blst, \
                     blhst, dpst, dphst, \
                     dist, lrt, yt = graph.get_collection("input_nodes")
@@ -114,7 +114,7 @@ def get_prediction(smg, sml, indexes, bls, blhs, dis, dps, dphs, model_name):
             preds_t_2 = sess.run([predictions], feed_dict=fd)
             preds_t_2 = np.array(preds_t_2)[0]
             preds = preds + w[i - 1] * (preds_t_1 + preds_t_2) / 2
-        tf.reset_default_graph()
+        tf.compat.v1.reset_default_graph()
     preds = np.argmax(preds, axis=1)
     print(preds.shape)
     return preds