diff --git a/code/project/code/autoscaled-inference-service.yaml b/code/project/code/autoscaled-inference-service.yaml
index 4c88856..e7e422d 100644
--- a/code/project/code/autoscaled-inference-service.yaml
+++ b/code/project/code/autoscaled-inference-service.yaml
@@ -11,5 +11,6 @@ spec:
     model:
       modelFormat:
         name: tensorflow
+      # This is only needed on Mac M1
       image: "emacski/tensorflow-serving:2.6.0"
       storageUri: "pvc://strategy-volume/saved_model_versions"
diff --git a/code/project/code/inference-service.yaml b/code/project/code/inference-service.yaml
index e386e5d..c2932fa 100644
--- a/code/project/code/inference-service.yaml
+++ b/code/project/code/inference-service.yaml
@@ -8,6 +8,7 @@ spec:
     model:
       modelFormat:
         name: tensorflow
+      # This is only needed on Mac M1
       image: "emacski/tensorflow-serving:2.6.0"
       # https://kserve.github.io/website/modelserving/storage/pvc/pvc/
       # Note that we are skipping `mountPath: /trained_model`
diff --git a/code/project/code/model-selection.py b/code/project/code/model-selection.py
index 312c17e..06db238 100644
--- a/code/project/code/model-selection.py
+++ b/code/project/code/model-selection.py
@@ -6,9 +6,6 @@
 import os
 
 
-BUFFER_SIZE = 10000
-BATCH_SIZE = 64
-
 # Scaling MNIST data from (0, 255] to (0., 1.]
 def scale(image, label):
   image = tf.cast(image, tf.float32)
@@ -21,7 +18,7 @@ def scale(image, label):
   model_path = "trained_model/saved_model_versions/" + str(i)
   model = keras.models.load_model(model_path)
   datasets, _ = tfds.load(name='fashion_mnist', with_info=True, as_supervised=True)
-  ds = datasets['test'].map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
+  ds = datasets['test'].map(scale).cache().shuffle(10000).batch(64)
   _, accuracy = model.evaluate(ds)
   if accuracy > best_accuracy:
     best_accuracy = accuracy
diff --git a/code/project/code/multi-worker-distributed-training.py b/code/project/code/multi-worker-distributed-training.py
index 9f09a79..d9f9ccc 100644
--- a/code/project/code/multi-worker-distributed-training.py
+++ b/code/project/code/multi-worker-distributed-training.py
@@ -23,7 +23,6 @@ def scale(image, label):
   return datasets['train'].map(scale).cache().shuffle(BUFFER_SIZE)
 
 
-# TODO: Use different models and pick top two for model serving
 def build_and_compile_cnn_model():
   print("Training CNN model")
   model = models.Sequential()
@@ -215,8 +214,6 @@ def is_chief():
   tf_config = json.loads(os.environ.get('TF_CONFIG') or '{}')
   TASK_INDEX = tf_config['task']['index']
 
-  # TODO: Add flag to train different models
-
   parser = argparse.ArgumentParser()
   parser.add_argument('--saved_model_dir',
                       type=str,
diff --git a/code/project/code/predict-service.py b/code/project/code/predict-service.py
index 3129852..d682c5a 100644
--- a/code/project/code/predict-service.py
+++ b/code/project/code/predict-service.py
@@ -6,9 +6,6 @@
 
 model = keras.models.load_model("trained_model/saved_model_versions")
 
-BUFFER_SIZE = 10000
-BATCH_SIZE = 64
-
 # Scaling MNIST data from (0, 255] to (0., 1.]
 def scale(image, label):
   image = tf.cast(image, tf.float32)
@@ -17,7 +14,7 @@ def scale(image, label):
 
 datasets, _ = tfds.load(name='fashion_mnist', with_info=True, as_supervised=True)
 
-ds = datasets['test'].map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
+ds = datasets['test'].map(scale).cache().shuffle(10000).batch(64)
 
 # TODO: Visualize the images and compare with the classified result
 model.predict(ds)
diff --git a/code/project/code/workflow.yaml b/code/project/code/workflow.yaml
index 2e0306c..248132e 100644
--- a/code/project/code/workflow.yaml
+++ b/code/project/code/workflow.yaml
@@ -5,13 +5,6 @@ metadata:
   namespace: kubeflow
 spec:
   entrypoint: tfjob-wf
-  podGC:
-    strategy: OnPodSuccess
-  volumes:
-  - name: model
-    persistentVolumeClaim:
-      claimName: strategy-volume
-
   templates:
   - name: tfjob-wf
     steps:
@@ -23,6 +16,12 @@ spec:
         template: model-selection-step
     - - name: create-model-serving-service
         template: create-model-serving-service
+  podGC:
+    strategy: OnPodSuccess
+  volumes:
+  - name: model
+    persistentVolumeClaim:
+      claimName: strategy-volume
 
   - name: data-ingestion-step
     serviceAccountName: argo