diff --git a/code/update-sdk.sh b/code/update-sdk.sh
index 4850e906..fe2cd7a3 100755
--- a/code/update-sdk.sh
+++ b/code/update-sdk.sh
@@ -1,12 +1,16 @@
 # Script that pulls the latest version of the SDK into your ThemeBuilder
 # To be used if there are any issues arise with pulling the latest SDK in your ThemeBuilder build
-CUR_DIR=`pwd`
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+BRANCH=main
+if [ $# -gt 0 ]; then
+   BRANCH=$1
+fi
 SDK=a11y-theme-builder-sdk
-cd ../../$SDK
-git checkout main
+cd $SCRIPT_DIR/../../$SDK
+git checkout $BRANCH
 git pull --rebase --autostash
 npm run build-lib
-cd $CUR_DIR
+cd $SCRIPT_DIR
 rm -rf src/ui/node_modules/.cache
 rm -rf src/ui/node_modules/$SDK
 cp -R ../../$SDK src/ui/node_modules