diff --git a/llvm/projects/keras/frontend/promise_translator.py b/llvm/projects/keras/frontend/promise_translator.py
index c12df7815677c30e7d7298284b5b8a6ed40dc024..54982bd7c7f906d5fe27ddb052ae303575c260ba 100644
--- a/llvm/projects/keras/frontend/promise_translator.py
+++ b/llvm/projects/keras/frontend/promise_translator.py
@@ -1038,7 +1038,7 @@ class PromiseRtTranslator:
       layer_ranges[layer_name] = []
       layer_it += 1
 
-    batch_size = 100
+    batch_size = 1000
     #batch_size = len(x_test)
     input_size = len(x_test)
     num_batches = input_size // batch_size 
@@ -1053,10 +1053,11 @@ class PromiseRtTranslator:
       end = (i + 1) * batch_size
 
       print ("start = ", start, " end = , ", end)
-
       
       # Inference over test set
       layer_outs = functor([x_test[start:end], 1.])
+
+      ###layer_outs = functor([x_test[start:end]])
     
       # NOTE: Saving quant ranges for input
       min_val = np.amin(x_test)
diff --git a/llvm/projects/keras/src/Benchmark.py b/llvm/projects/keras/src/Benchmark.py
index f37e4de6bcbdccf5d45d5570540c6c92e1c900a5..0d51cbb7e0c1c6434f17c644e6d25d9609140afb 100644
--- a/llvm/projects/keras/src/Benchmark.py
+++ b/llvm/projects/keras/src/Benchmark.py
@@ -1,11 +1,24 @@
 
 
+import sys
+from keras.utils.np_utils import to_categorical
+from keras.models import load_model
+from frontend.approxhpvm_translator import translate_to_approxhpvm
+from frontend.weight_utils import dumpCalibrationData
+from frontend.weight_utils import dumpHPVMToKerasModel
 
-class Benchmark:
 
-    def __init__(self):
-        return
 
+class Benchmark:
+
+    def __init__(self, name, reload_dir, keras_model_file, hpvm_dir, num_classes):
+        self.name = name
+        self.reload_dir = reload_dir
+        self.keras_model_file = keras_model_file
+        self.hpvm_dir = hpvm_dir
+        self.num_classes = num_classes
+        
+        
     def buildModel(self):
         return
 
@@ -18,8 +31,36 @@ class Benchmark:
     def inference(self):
         return
 
-    def run(self):
-        return
 
-        
+    def run(self, argv):
+
+      if len(argv) < 2:
+        sys.exit(0)
+
+      model = self.buildModel()
+
+      X_train, Y_train, X_test, Y_test = self.data_preprocess()   
+
+      if argv[1] == "hpvm_reload":
+        model = dumpHPVMToKerasModel(model, self.reload_dir, self.keras_model_file, X_test, Y_test)
+
+      if argv[1] == "keras_reload":
+        model = load_model(self.keras_model_file)
+
+      if argv[1] == "train":
+        model = self.trainModel(model)
+
+      score = model.evaluate(X_test, to_categorical(Y_test, self.num_classes), verbose=0)
+      print('Test accuracy2:', score[1])
+
+
+      if len(argv) > 2 and argv[2] == "frontend":
+        if argv[1] == "keras_reload":
+          print("ERROR: Must load HPVM model to invoke frontend")
+          sys.exit(1)
+
+        translate_to_approxhpvm(model, self.hpvm_dir, X_test, Y_test, self.num_classes)
+
+
+
 
diff --git a/llvm/projects/keras/src/alexnet.py b/llvm/projects/keras/src/alexnet.py
index d10bfa6b9b7f5966fe541c04c1260f7a4d0b3aa9..756af48b7cb6f338171df59a5c8bd399cd248e4b 100644
--- a/llvm/projects/keras/src/alexnet.py
+++ b/llvm/projects/keras/src/alexnet.py
@@ -28,8 +28,8 @@ from frontend.weight_utils import dumpHPVMToKerasModel
 class AlexNet(Benchmark):
 
 
-  def __init__(self):
-      self.name = "AlexNet"
+  #def __init__(self):
+  #    self.name = "AlexNet"
       
       
   def lr_schedule(self, epoch):
@@ -193,46 +193,25 @@ class AlexNet(Benchmark):
     return X_train, Y_train, X_test, Y_test
 
 
-
+  
 
     
 if __name__ == "__main__":
 
-    if len(sys.argv) < 2:
-      sys.exit(0)
-    
+      
     os.environ["CUDA_VISIBLE_DEVICES"] = "0"
     # Changing to NCHW format
     K.set_image_data_format('channels_first')
 
-    alexnet = AlexNet()
-    
-    model = alexnet.buildModel()
 
-    X_train, Y_train, X_test, Y_test = alexnet.data_preprocess()   
+    ### Parameters specific to each benchmark
     reload_dir = "/home/hsharif3/Gitlab/hpvm/llvm/projects/hpvm-tensor-rt/model_params/alexnet_cifar10/"
     keras_model_file = "alexnet.h5"
-
-    if sys.argv[1] == "hpvm_reload":
-      model = dumpHPVMToKerasModel(model, reload_dir, keras_model_file, X_test, Y_test)
-
-    if sys.argv[1] == "keras_reload":
-      model = load_model(keras_model_file)
- 
-    if sys.argv[1] == "train":
-      model = alexnet.trainModel(model)
-
+    hpvm_dir = "data/alexnet_cifar10/" 
     num_classes = 10
-    score = model.evaluate(X_test, to_categorical(Y_test, num_classes), verbose=0)
-    print('Test accuracy2:', score[1])
 
-      
-    if len(sys.argv) > 2 and sys.argv[2] == "frontend":
-      if sys.argv[1] == "keras_reload":
-        print("ERROR: Must load HPVM model to invoke frontend")
-        sys.exit(1)
-        
-      hpvm_dir = "data/alexnet_cifar10/" 
-      translate_to_approxhpvm(model, hpvm_dir, X_test, Y_test, num_classes)
+    alexnet = AlexNet("AlexNet", reload_dir, keras_model_file, hpvm_dir, num_classes)
     
+    alexnet.run(sys.argv)
+