diff --git a/.github/workflows/pypi_deployment.yml b/.github/workflows/pypi_deployment.yml
index fbad5f3..e128e95 100644
--- a/.github/workflows/pypi_deployment.yml
+++ b/.github/workflows/pypi_deployment.yml
@@ -94,7 +94,7 @@ jobs:
 
     - uses: actions/checkout@v4
       with:
-        ref: "main"
+        ref: "dev-${{ env.VERSION }}"
 
     - name: Install dependencies
       run: |
diff --git a/scripts/pomdp_py_release.sh b/scripts/pomdp_py_release.sh
index dfc11c1..1d2c4f5 100755
--- a/scripts/pomdp_py_release.sh
+++ b/scripts/pomdp_py_release.sh
@@ -87,10 +87,8 @@ fi
 echo $version
 if ! is_git_repo_on_branch $pomdp_py_path dev-$version; then
     if ! is_git_repo_on_branch $pomdp_py_path dev-latest; then
-        if ! is_git_repo_on_branch $pomdp_py_path main; then
-            echo "pomdp-py repo must be either on dev-latest or dev-$version, or main, but not $current_branch. Abort"
-            exit 1
-        fi
+        echo "pomdp-py repo must be either on dev-latest or dev-$version, but not $current_branch. Abort"
+        exit 1
     fi
 fi