diff --git a/.azure-pipelines/template/ut-template.yml b/.azure-pipelines/template/ut-template.yml
index f31eecbe..54f10966 100644
--- a/.azure-pipelines/template/ut-template.yml
+++ b/.azure-pipelines/template/ut-template.yml
@@ -35,16 +35,14 @@ steps:
   - ${{ if eq(parameters.imageSource, 'build') }}:
     - script: |
         docker exec ${{ parameters.utContainerName }} bash -c "cd /auto-round \
-          && python setup.py bdist_wheel \
-          && pip install $(ls dist/*.whl | head -n 1)[cpu] \
+          && pip install .[cpu] \
           && pip list"
       displayName: "Env Setup"
 
   - ${{ if eq(parameters.imageSource, 'pull') }}:
     - script: |
         docker exec ${{ parameters.utContainerName }} bash -c "cd /auto-round \
-          && python setup.py bdist_wheel \
-          && pip install $(ls dist/*.whl | head -n 1)[hpu] \
+          && pip install .[hpu] \
           && pip list"
       displayName: "HPU Env Setup"