Using find_library to verify Torchvision package

find_library() function in CMake is used to search for a library in specific directories and return its path if found.

Co-authored-by: Misha Chornyi <[email protected]>

Author: pvijayakrish

CMakeLists.txt

@@ -78,16 +78,10 @@ else()
   endif()
 
     # Look for installed Torchvision package in lib paths
-  if(TRITON_PYTORCH_ENABLE_TORCHVISION)
-    if(${RHEL_BUILD})
-      if(NOT EXISTS "${TRITON_PYTORCH_LIB_PATHS}/libtorchvision.so")
-        message(WARNING "TRITON_PYTORCH_ENABLE_TORCHVISION is on, but TRITON_PYTORCH_LIB_PATHS does not contain Torchvision package")
-      endif()
-    else()
-      if(NOT EXISTS "${TRITON_PYTORCH_LIB_PATHS}/libtorchvision.so.1")
-        message(WARNING "TRITON_PYTORCH_ENABLE_TORCHVISION is on, but TRITON_PYTORCH_LIB_PATHS does not contain Torchvision package")
-      endif()
-    endif()
+  find_library( LIBTORCHVISION libtorchvision.so libtorchvision.so.1 PATHS ${TRITON_PYTORCH_LIB_PATHS} )
+  if(NOT ${LIBTORCHVISION}) 
+    message(WARNING "TRITON_PYTORCH_ENABLE_TORCHVISION is on, but TRITON_PYTORCH_LIB_PATHS does not contain Torchvision package")
+  endif(NOT ${LIBTORCHVISION})
   endif()
 endif()