Add option to expose custom PyTorch repo/branch (llvm#1103)

tanyokwok · Jul 25, 2022 · f424930 · f424930
1 parent b80ce79
commit f424930
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -65,6 +65,8 @@ option(TORCH_MLIR_USE_INSTALLED_PYTORCH "Build from local PyTorch in environment
 if(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER)
   if (NOT TORCH_MLIR_USE_INSTALLED_PYTORCH)
     # Source builds
+    set(ENV{PYTORCH_REPO} ${PYTORCH_REPO})
+    set(ENV{PYTORCH_BRANCH} ${PYTORCH_BRANCH})
     set(ENV{MACOSX_DEPLOYMENT_TARGET} ${MACOSX_DEPLOYMENT_TARGET})
     set(ENV{CMAKE_OSX_ARCHITECTURES} ${CMAKE_OSX_ARCHITECTURES})
     set(ENV{CMAKE_C_COMPILER_LAUNCHER} ${CMAKE_C_COMPILER_LAUNCHER})