diff --git a/cli/setup.sh b/cli/setup.sh
index 8914f24309..0e23cc7dee 100644
--- a/cli/setup.sh
+++ b/cli/setup.sh
@@ -20,7 +20,8 @@
 
 
 # <az_ml_install>
-az extension add -n ml -y
+# az extension add -n ml -y
+az extension add --source https://azuremlsdktestpypi.blob.core.windows.net/wheels/sdk-cli-v2-public/ml-2.34.0-py3-none-any.whl --yes
 # </az_ml_install>
 
 ## For backward compatibility - running on old subscription
diff --git a/sdk/python/setup.sh b/sdk/python/setup.sh
index ec29a18244..86de911c3a 100644
--- a/sdk/python/setup.sh
+++ b/sdk/python/setup.sh
@@ -1,7 +1,6 @@
 #!/bin/bash
 
 # <az_ml_sdk_install>
-# pip install --pre azure-ai-ml
 # </az_ml_sdk_install>
 
 # <mldesigner_install>
@@ -15,9 +14,7 @@ pip install pandas
 
 
 # <az_ml_sdk_test_install>
-# pip install azure-ai-ml==0.1.0.b8
-pip install azure-ai-ml
-# https://docsupport.blob.core.windows.net/ml-sample-submissions/1905732/azure_ai_ml-1.0.0-py3-none-any.whl
+pip install https://docsupport.blob.core.windows.net/ml-sample-submissions/4446622/azure_ai_ml-1.24.0-py3-none-any.whl
 # </az_ml_sdk_test_install>
 
 # protobuf==5.29.0 has IndentationError bug