diff --git a/coreml.py b/coreml.py
index db2fbae..9c091d8 100644
--- a/coreml.py
+++ b/coreml.py
@@ -1,8 +1,8 @@
 import coremltools as ct
 
 
-def convert_coreml(model_path, save_path='../mobilenet_nnconv5.mlmodel'):
-    mlmodel = ct.convert(model_path)
+def convert_coreml(model_path, save_path='mobilenet_nnconv5.mlmodel'):
+    mlmodel = ct.convert(model_path, inputs=[ct.ImageType()])
     mlmodel.save(save_path)
 
 
diff --git a/fast_depth_functional.py b/fast_depth_functional.py
index 672f224..91eb943 100644
--- a/fast_depth_functional.py
+++ b/fast_depth_functional.py
@@ -72,6 +72,7 @@ def mobilenet_nnconv5(weights=None, shape=(224, 224, 3)):
     x = keras.layers.Conv2D(1, 1, padding='same')(x)
     x = keras.layers.BatchNormalization()(x)
     x = keras.layers.ReLU(6.)(x)
+    x = keras.layers.Reshape([shape[0], shape[1]])(x)
     return keras.Model(inputs=input, outputs=x, name="fast_depth")