diff --git a/backend/MetaAugment/Evo_learner.py b/backend/MetaAugment/Evo_learner.py
index 71170f437b7674a15e8941d0a15c6990595ab55d..d111a9a8b5e726734fe9fb971402ed5cc8f28eef 100644
--- a/backend/MetaAugment/Evo_learner.py
+++ b/backend/MetaAugment/Evo_learner.py
@@ -19,6 +19,7 @@ from typing import List, Tuple, Optional, Dict
 import heapq
 import math
 import torch
+from flask import current_app
 
 from enum import Enum
 from torch import Tensor
@@ -94,6 +95,7 @@ class Evolutionary_learner():
         self.sub_num_pol = sub_num_pol
         self.mag_bins = mag_bins
         self.fun_num = fun_num
+        self.iter_count = 0
 
         full_augmentation_space = [
             # (function_name, do_we_need_to_specify_magnitude)
@@ -317,6 +319,9 @@ class Evolutionary_learner():
             fit_val = ((test_autoaugment_policy(full_policy, self.train_dataset, self.test_dataset, self.child_network)[0])/
                         + test_autoaugment_policy(full_policy, self.train_dataset, self.test_dataset, self.child_network)[0]) / 2
 
+            self.iter_count += 1
+            current_app.config['iteration'] = self.iter_count
+
             return fit_val
 
         def on_generation(ga_instance):
diff --git a/backend/app.py b/backend/app.py
index df6e34a23814bd09d79ca969f543fb49a4afc1ef..5e39517f6ae17dc93910e02f960e7aac0074dd7d 100644
--- a/backend/app.py
+++ b/backend/app.py
@@ -2,8 +2,9 @@
 # if __name__ == '__main__':
 #     app.run(host='0.0.0.0',port=port)
 
+from numpy import broadcast
 from auto_augmentation import home, progress,result, training
-
+from flask_socketio import SocketIO,  send
 
 from flask import Flask, flash, request, redirect, url_for
 from werkzeug.utils import secure_filename
@@ -14,12 +15,10 @@ app = create_app()
 
 
 
-
 # app = Flask(__name__)
 # app.config.from_mapping(
 #     SECRET_KEY='dev',
 # )
-# # ensure the instance folder exists
 # os.makedirs(app.instance_path, exist_ok=True)
 # from auto_augmentation import download_file
 # app.register_blueprint(home.bp)
@@ -30,10 +29,17 @@ app = create_app()
 
 
 
-# port = int(os.environ.get("PORT", 5000))
+socketio = SocketIO(app)
+
+
 
+@socketio.on('message')
+def handleMessage(msg):
+    print("Message: ", msg)
+    send(msg, broadcast=True)
 
 
 
 if __name__ == '__main__':
-    app.run(debug=False)
\ No newline at end of file
+    app.run(debug=True)
+    # socketio.run(app)
\ No newline at end of file
diff --git a/backend/auto_augmentation/progress.py b/backend/auto_augmentation/progress.py
index d1bc33384d82d354e0812fbcd82e3b762139c2d3..d175dce8375a9490ee45c61abd3abda23a6dce28 100644
--- a/backend/auto_augmentation/progress.py
+++ b/backend/auto_augmentation/progress.py
@@ -51,8 +51,6 @@ def response():
 
         # if user upload datasets and networks, save them in the database
 
-        # print("DS HERE3: ", ds)
-        # print("ds up: ", ds_up)
         if ds == None and ds_up != None:
             ds = 'Other'
             ds_folder = request.files['dataset_upload']
diff --git a/backend/auto_augmentation/templates/progress.html b/backend/auto_augmentation/templates/progress.html
index d0a6a8bdd137a1d593be5dfcf06fd68eebcd1f24..39ace7b0b7ae89aba789306e6e8ef3c4fb99c723 100644
--- a/backend/auto_augmentation/templates/progress.html
+++ b/backend/auto_augmentation/templates/progress.html
@@ -5,7 +5,7 @@
 
 <br>
 <ul>
-    Training the {{auto_aug_leanrer}} agent...
+    Training the {{auto_aug_learner}} agent...
 
 </ul>
 
@@ -41,9 +41,39 @@
     }
     </script>
 
+    <!-- <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.3.1/jquery.min.js"></script> -->
+    <!-- <script src="https://cdnjs.cloudflare.com/ajax/libs/socket.io/4.0.1/socket.io.js" integrity="sha512-q/dWJ3kcmjBLU4Qc47E4A9kTB4m3wuTY7vkFJDTZKjTs8jhyGQnaUrxa0Ytd0ssMZhbNua9hE+E7Qv1j+DyZwA==" crossorigin="anonymous"></script> -->
+    <!-- <script type="text/javascript" charset="utf-8"></script> -->
+
+
     <form action="/show_result">
         <input type="submit" value='Show Result'>
     </form>
+
+
+    <script type='text/javascript'>
+
+        $(document).ready(function() {
+        var socket  = io.connect('http://127.0.0.1:5000');
+
+        socket.on('connect', function() {
+            socket.send('User has connected');
+        });
+
+        socket.on('message', function(msg) {
+             $('#messages').append('<li>'+msg+'</li>');
+             console.log('Recieved message');
+        });
+
+        });      
+
+
+    </script>
+
+    <ul id='messages'></ul>
+
+
+
 </div>
  
 {% endblock %}
\ No newline at end of file
diff --git a/backend/auto_augmentation/training.py b/backend/auto_augmentation/training.py
index 222548ceb6b22377b1d79892a0df4b026474f436..39001e9eb3bdbada40daef5cc443007e342206ac 100644
--- a/backend/auto_augmentation/training.py
+++ b/backend/auto_augmentation/training.py
@@ -53,8 +53,6 @@ def response():
     if auto_aug_learner == 'UCB':
         policies = UCB1_JC.generate_policies(num_policies, num_sub_policies)
         q_values, best_q_values = UCB1_JC.run_UCB1(policies, batch_size, learning_rate, ds, toy_size, max_epochs, early_stop_num, iterations, IsLeNet, ds_name)     
-        # plt.figure()
-        # plt.plot(q_values)
         best_q_values = np.array(best_q_values)
 
     elif auto_aug_learner == 'Evolutionary Learner':