diff --git a/.github/workflows/run-setup-script.yml b/.github/workflows/run-setup-script.yml
index bfba94b..c41727a 100644
--- a/.github/workflows/run-setup-script.yml
+++ b/.github/workflows/run-setup-script.yml
@@ -2,17 +2,24 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-on: [push, pull_request]
 name: Azure ARM
+on:
+  push:
+  # Runs tests for pushes to main branch.
+    branches:
+      - main
+  # Runs tests for any PRs.
+  pull_request:
+    branches:
+      - main
 jobs:
   build-and-deploy:
     runs-on: ubuntu-latest
     steps:
-
-      # Checkout code
+    # Checkout code
     - uses: actions/checkout@main
 
-      # Log into Azure
+    # Log into Azure
     - uses: azure/login@v1
       with:
         creds: ${{ secrets.AZURE_CREDENTIALS }}
@@ -30,9 +37,9 @@ jobs:
 
     - name: Grant Admin role for synapse workspace to github service-principal and AAD group
       run: |
-        # Granting role to principal-id used in github workflow
-        az synapse role assignment create --workspace-name pr${{ steps.vars.outputs.sha_short }}-pipeline-syn-ws --role "Synapse Administrator" --assignee $(${{ secrets.AZURE_CREDENTIALS }} | jq '.clientId')
-        # Granting role to principal-id used in AAD group
+        # Granting role to principal-id for github workflow
+        az synapse role assignment create --workspace-name pr${{ steps.vars.outputs.sha_short }}-pipeline-syn-ws --role "Synapse Administrator" --assignee $(az account show --query 'user.name' -o tsv)
+        # Granting role to principal-id for AAD group
         az synapse role assignment create --workspace-name pr${{ steps.vars.outputs.sha_short }}-pipeline-syn-ws --role "Synapse Administrator" --assignee ${{ secrets.AAD_GROUP_PRINCIPAL_ID }}
       shell: bash
 
@@ -40,6 +47,25 @@ jobs:
       run: BATCH_ACCOUNT_NAME=${{ secrets.CI_BATCH_ACCOUNT_NAME }} SOURCE_STORAGE_ACCOUNT_NAME=${{ secrets.CI_STORAGE_ACCOUNT }} ./test/deploy-prepare-and-test-pipeline.sh pr${{ steps.vars.outputs.sha_short }}
       shell: bash
 
+  cleanup:
+    runs-on: ubuntu-latest
+    needs: build-and-deploy
+    if: always()
+    steps:
+    # Checkout code
+    - uses: actions/checkout@main
+    # Log into Azure
+    - uses: azure/login@v1
+      with:
+        creds: ${{ secrets.AZURE_CREDENTIALS }}
+
+    - name: Get git info
+      id: vars
+      shell: bash
+      run: |
+        echo "##[set-output name=branch;]$(echo ${GITHUB_REF#refs/heads/})"
+        echo "::set-output name=sha_short::$(git rev-parse --short HEAD)"
+
     - name: Cleanup the resources
-      run: ./deploy/cleanup.sh pr${{ steps.vars.outputs.sha_short }} ${{ secrets.CI_BATCH_ACCOUNT_NAME }} custom-vision-model-v2
+      run: ./deploy/cleanup.sh pr${{ steps.vars.outputs.sha_short }} ${{ secrets.CI_BATCH_ACCOUNT_NAME }}
       shell: bash
\ No newline at end of file