diff --git a/llvm/projects/keras/src/mobilenet_shallow.py b/llvm/projects/keras/src/mobilenet_shallow.py
index 442004ef8e4148bf70e562a90784c2bcda5d753e..ca5d78d84144dd663132eb173393b027d81c9ce0 100644
--- a/llvm/projects/keras/src/mobilenet_shallow.py
+++ b/llvm/projects/keras/src/mobilenet_shallow.py
@@ -1,4 +1,3 @@
-
 import sys
 import os
 os.environ['CUDA_VISIBLE_DEVICES'] = '0'
@@ -42,59 +41,68 @@ def get_mobilenet(alpha=1, depth_multiplier=1):
     model = Sequential()
     
     def _conv_block(filters, alpha, kernel=(3, 3), strides=(1, 1)):
-        channel_axis = 1
-        
+        channel_axis = 1 if K.image_data_format() == 'channels_first' else -1
+        filters = int(filters * alpha)
         model.add(Conv2D(filters, kernel,
                           padding='same',
                           use_bias=False,
                           strides=strides, 
-                          input_shape=(3, 32, 32)))
+                        input_shape=(32, 32, 3)))
         model.add(BatchNormalization(axis=channel_axis))
-        model.add(Activation('relu'))
+        model.add(ReLU())
     
     def _depthwise_conv_block(pointwise_conv_filters, alpha, depth_multiplier=1, strides=(1, 1)):
-        channel_axis = 1 
-
-        model.add(ZeroPadding2D(padding = ((1,1), (1,1) )))
+        channel_axis = 1 if K.image_data_format() == 'channels_first' else -1
+        pointwise_conv_filters = int(pointwise_conv_filters * alpha)
 
         model.add(DepthwiseConv2D((3, 3),
-                                   padding='valid',
-                                   #depth_multiplier=depth_multiplier,
+                                   padding='same',
+                                   depth_multiplier=depth_multiplier,
                                    strides=strides,
                                    use_bias=False))    
         model.add(BatchNormalization(axis=channel_axis))
-        
-        model.add(Activation('relu'))
+        model.add(ReLU())
         model.add(Conv2D(pointwise_conv_filters, (1, 1),
                           padding='same',
                           use_bias=False,
                           strides=(1, 1)))
         model.add(BatchNormalization(axis=channel_axis))
-        model.add(Activation('relu'))
-
+        model.add(ReLU())
+        
 
     _conv_block(32, alpha, strides=(1, 1))
     
     _depthwise_conv_block(64, alpha, depth_multiplier)
-
-    model.add(Dropout(rate=0.4))
-
-    _depthwise_conv_block(64, alpha, depth_multiplier,
+    
+    _depthwise_conv_block(128, alpha, depth_multiplier,
                               strides=(2, 2))
-    _depthwise_conv_block(128, alpha, depth_multiplier, strides=(2, 2))
-    model.add(Dropout(rate=0.4))
+    _depthwise_conv_block(128, alpha, depth_multiplier)
+    model.add(Dropout(rate=0.5))
 
     _depthwise_conv_block(256, alpha, depth_multiplier, 
                       strides=(2, 2))
     _depthwise_conv_block(256, alpha, depth_multiplier)
-    model.add(Dropout(rate=0.4))
+    model.add(Dropout(rate=0.5))
+
+    _depthwise_conv_block(512, alpha, depth_multiplier,
+                      strides=(2, 2))
+#     _depthwise_conv_block(512, alpha, depth_multiplier)
+#     _depthwise_conv_block(512, alpha, depth_multiplier)
+#     model.add(Dropout(rate=0.5))
+    
+#     _depthwise_conv_block(512, alpha, depth_multiplier)
+#     _depthwise_conv_block(512, alpha, depth_multiplier)
+#     _depthwise_conv_block(512, alpha, depth_multiplier)
+#     model.add(Dropout(rate=0.5))
+    
+#     _depthwise_conv_block(1024, alpha, depth_multiplier,
+#                          strides=(2, 2))
+#     _depthwise_conv_block(1024, alpha, depth_multiplier)
+#     model.add(Dropout(rate=0.5))
 
     model.add(AveragePooling2D(pool_size=2))
     model.add(Flatten())
-    model.add(Dense(10))
-    #, activation='softmax'))
-    model.add(Activation('softmax'))
-
+    model.add(Dense(10, activation='softmax'))
 
     return model