Do not link _pywrap_tensorflow_internal.so if XLA is not enabled.

Signed-off-by: Trent Lo <trentl@nvidia.com>
horovod · Aug 12, 2021 · 1a031b3 · 1a031b3
1 parent 0b7e04d
commit 1a031b3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/cmake/Modules/FindTensorflow.cmake b/cmake/Modules/FindTensorflow.cmake
@@ -21,8 +21,10 @@ if (LEN EQUAL "4")
     list(GET Tensorflow_OUTPUT 2 Tensorflow_LIBRARIES)
     string(REPLACE " " ";" Tensorflow_LIBRARIES_LIST "${Tensorflow_LIBRARIES}")
     list(GET Tensorflow_LIBRARIES_LIST 0 Tensorflow_LIB_PATH)
-    # XLA implementations are in _pywrap_tensorflow_internal.so
-    set(Tensorflow_LIBRARIES "${Tensorflow_LIBRARIES} ${Tensorflow_LIB_PATH}/python/ -l:_pywrap_tensorflow_internal.so")
+    if (Tensorflow_VERSION VERSION_GREATER_EQUAL "2.6")
+        # XLA implementations are in _pywrap_tensorflow_internal.so
+        set(Tensorflow_LIBRARIES "${Tensorflow_LIBRARIES} ${Tensorflow_LIB_PATH}/python/ -l:_pywrap_tensorflow_internal.so")
+    endif()
     message("Tensorflow_LIBRARIES := ${Tensorflow_LIBRARIES}")
     list(GET Tensorflow_OUTPUT 3 Tensorflow_COMPILE_FLAGS)
     if("${Tensorflow_COMPILE_FLAGS}" MATCHES "-D_GLIBCXX_USE_CXX11_ABI=1")