diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
index b7a807d264d..043fadc3538 100644
--- a/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
+++ b/test/cpp_extensions/open_registration_extension/torch_openreg/cmake/TorchPythonTargets.cmake
@@ -1,9 +1,7 @@
 if(WIN32)
-  set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/torch_python.lib")
-elseif(APPLE)
-  set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/libtorch_python.dylib")
+  set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/${CMAKE_IMPORT_LIBRARY_PREFIX}torch_python${CMAKE_IMPORT_LIBRARY_SUFFIX}")
 else()
-  set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/libtorch_python.so")
+  set(TORCH_PYTHON_IMPORTED_LOCATION "${PYTORCH_INSTALL_DIR}/lib/${CMAKE_SHARED_LIBRARY_PREFIX}torch_python${CMAKE_SHARED_LIBRARY_SUFFIX}")
 endif()
 
 add_library(torch_python SHARED IMPORTED)