diff --git a/cli/setup.sh b/cli/setup.sh
index 8914f24309..1786477341 100644
--- a/cli/setup.sh
+++ b/cli/setup.sh
@@ -20,7 +20,7 @@
 
 
 # <az_ml_install>
-az extension add -n ml -y
+az extension add --source https://azuremlsdktestpypi.blob.core.windows.net/wheels/sdk-cli-v2-public/ml-2.32.0-py3-none-any.whl --upgrade -y
 # </az_ml_install>
 
 ## For backward compatibility - running on old subscription