diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
index c74c40a21..51b154932 100644
--- 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})