Do not install upstream pytorch when running benchmarks (#2118)

Signed-off-by: Whitney Tsang <[email protected]>
intel · Sep 4, 2024 · 02f5eed · 02f5eed
1 parent cb0b7e1
commit 02f5eed
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/scripts/test-triton.sh b/scripts/test-triton.sh
@@ -122,7 +122,11 @@ if [ "$TRITON_TEST_REPORTS" == true ]; then
     capture_runtime_env
 fi
 
-$SKIP_DEPS || $SCRIPTS_DIR/compile-pytorch-ipex.sh --upstream-pytorch --pinned $([ $VENV = true ] && echo "--venv")
+if [ "$TEST_BENCHMARK_SOFTMAX" = true ] || [ "$TEST_BENCHMARK_GEMM" = true ] || [ "$TEST_BENCHMARK_ATTENTION" = true ]; then
+  $SKIP_DEPS || $SCRIPTS_DIR/compile-pytorch-ipex.sh --pytorch --ipex --pinned --source $([ $VENV = true ] && echo "--venv")
+else
+  $SKIP_DEPS || $SCRIPTS_DIR/compile-pytorch-ipex.sh --upstream-pytorch --pinned $([ $VENV = true ] && echo "--venv")
+fi
 
 if [ ! -d "$TRITON_PROJ_BUILD" ]
 then
@@ -317,9 +321,6 @@ test_triton() {
   if [ "$TEST_MICRO_BENCHMARKS" = true ]; then
     run_microbench_tests
   fi
-  if [ "$TEST_BENCHMARK_SOFTMAX" = true ] || [ "$TEST_BENCHMARK_GEMM" = true ] || [ "$TEST_BENCHMARK_ATTENTION" = true ]; then
-    $SKIP_DEPS || $SCRIPTS_DIR/compile-pytorch-ipex.sh --pinned --source --pytorch --ipex
-  fi
   if [ "$TEST_BENCHMARK_SOFTMAX" = true ]; then
     run_benchmark_softmax
   fi