diff --git a/downloadblob.sh b/downloadblob.sh
index af749ba..795cf62 100644
--- a/downloadblob.sh
+++ b/downloadblob.sh
@@ -8,9 +8,9 @@
 
 # Downloading Azure CLI on the VSTS build agent machine
 apt-get update -y && apt-get install -y python libssl-dev libffi-dev python-dev build-essential
-curl -L https://azurecliprod.blob.core.windows.net/install.py -o install.py
-printf "/usr/azure-cli\n/usr/bin" | python install.py
-az
+#curl -L https://azurecliprod.blob.core.windows.net/install.py -o install.py
+#printf "/usr/azure-cli\n/usr/bin" | python install.py
+#az
 
 #Setting environment variables to access the blob container
 export AZURE_STORAGE_ACCOUNT=$1
@@ -20,7 +20,9 @@ export blob_name1=$4
 export blob_name2=$5
 
 # Downloading Blob
-mkdir flaskwebapp
+if [ ! d flaskwebapp ]; then
+    mkdir flaskwebapp
+fi
 az storage blob download --container-name $container_name --name $blob_name1 --file flaskwebapp/$blob_name1 --output table
 az storage blob download --container-name $container_name --name $blob_name2 --file flaskwebapp/$blob_name2 --output table
 az storage blob list --container-name $container_name --output table
diff --git a/flaskwebapp/dockerfile b/flaskwebapp/dockerfile
index ee92044..f2ce86c 100644
--- a/flaskwebapp/dockerfile
+++ b/flaskwebapp/dockerfile
@@ -18,6 +18,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     rm /etc/nginx/sites-enabled/default && \
     cp /code/nginx/app /etc/nginx/sites-available/ && \
     ln -s /etc/nginx/sites-available/app /etc/nginx/sites-enabled/ && \
+    pip install wheel && \
     pip install -r /code/requirements.txt && \
     pip install https://cntk.ai/PythonWheel/CPU-Only/cntk-2.0-cp27-cp27mu-linux_x86_64.whl
 
diff --git a/flaskwebapp/driver.py b/flaskwebapp/driver.py
index b80e120..0b33baa 100644
--- a/flaskwebapp/driver.py
+++ b/flaskwebapp/driver.py
@@ -33,7 +33,7 @@ def init():
 
     # Load model and load the model from brainscript (3rd index)
     trainedModel = load_model(MODEL_FILE)
-    trainedModel = combine([trainedModel.outputs[3].owner])
+    trainedModel = combine([trainedModel.outputs[2].owner])
     end = t.default_timer()
 
     loadTimeMsg = "Model loading time: {0} ms".format(round((end - start) * 1000, 2))
diff --git a/flaskwebapp/requirements.txt b/flaskwebapp/requirements.txt
index dacae23..573c431 100644
--- a/flaskwebapp/requirements.txt
+++ b/flaskwebapp/requirements.txt
@@ -1,3 +1,4 @@
+setuptools >= 0.8.0
 pillow
 click==6.7
 configparser==3.5.0
diff --git a/tests/integration/test_api.sh b/tests/integration/test_api.sh
index 4e62529..df6f85d 100644
--- a/tests/integration/test_api.sh
+++ b/tests/integration/test_api.sh
@@ -6,10 +6,11 @@ MODEL_API_URL=$1
 # Simple API test
 echo "Testing API"
 reply=$(curl -s $MODEL_API_URL/version)
-if [[ $reply == 2.0rc1 ]]; then
+if [[ $reply == 2.0 ]]; then
         echo -e "Successfully validated version API call"
 else
         echo "Basic API call failed"
+        echo "Reply: "$reply
         exit 1
 fi