Update the RHEL library path

Author: pvijayakrish

CMakeLists.txt

@@ -78,8 +78,16 @@ else()
   endif()
 
     # Look for installed Torchvision package in lib paths
-  if(TRITON_PYTORCH_ENABLE_TORCHVISION AND 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")
+  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()
   endif()
 endif()
 
@@ -164,7 +172,7 @@ set(PT_LIBS
 if (${TRITON_PYTORCH_ENABLE_TORCHVISION})
   set(PT_LIBS
       ${PT_LIBS}
-      "libtorchvision.so.1"
+      ${RHEL_BUILD} ? "libtorchvision.so" : "libtorchvision.so.1"
   )
 endif() # TRITON_PYTORCH_ENABLE_TORCHVISION
 
@@ -249,7 +257,8 @@ if (${TRITON_PYTORCH_DOCKER_BUILD})
     COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH}/torch/lib/libcaffe2_nvrtc.so libcaffe2_nvrtc.so
     # TODO: Revisit when not needed by making it part of cuda base container.
     COMMAND docker cp -L pytorch_backend_ptlib:/usr/local/cuda/lib64/libcusparseLt.so libcusparseLt.so;
-    COMMAND /bin/sh -c "if [ ${TRITON_PYTORCH_ENABLE_TORCHVISION} = 'ON' ]; then docker cp  -a -L pytorch_backend_ptlib:/usr/local/${LIB_DIR}/libtorchvision.so.1 libtorchvision.so.1; fi"
+    #COMMAND /bin/sh -c "if [ ${TRITON_PYTORCH_ENABLE_TORCHVISION} = 'ON' ]; then docker cp  -a -L pytorch_backend_ptlib:/usr/local/${LIB_DIR}/libtorchvision.so.1 libtorchvision.so.1; fi"
+    COMMAND /bin/sh -c "if [ ${TRITON_PYTORCH_ENABLE_TORCHVISION} = 'ON' ]; then if [ ${RHEL_BUILD} = 'ON' ]; then docker cp -a -L pytorch_backend_ptlib:/usr/local/lib64/libtorchvision.so libtorchvision.so; else docker cp -a -L pytorch_backend_ptlib:/usr/local/${LIB_DIR}/libtorchvision.so.1 libtorchvision.so.1; fi; fi"
     COMMAND /bin/sh -c "if [ ${TRITON_PYTORCH_ENABLE_TORCHVISION} = 'ON' ]; then docker cp pytorch_backend_ptlib:/opt/pytorch/vision/torchvision/csrc include/torchvision/torchvision; fi"
     COMMAND /bin/sh -c "if [ ${TRITON_PYTORCH_ENABLE_TORCHTRT} = 'ON' ]; then docker cp pytorch_backend_ptlib:/usr/local/lib/python3.12/dist-packages/torch_tensorrt/lib/libtorchtrt_runtime.so libtorchtrt_runtime.so; fi"
     COMMAND docker cp pytorch_backend_ptlib:${PY_INSTALL_PATH}/torch_tensorrt/bin/torchtrtc torchtrtc || echo "error ignored..." || true
@@ -382,9 +391,15 @@ if (${TRITON_PYTORCH_DOCKER_BUILD})
   set(TRITON_PYTORCH_LIBS "${CMAKE_CURRENT_BINARY_DIR}/libtorch.so")
 
   if (${TRITON_PYTORCH_ENABLE_TORCHVISION})
-    set(TRITON_PYTORCH_LIBS
-        ${TRITON_PYTORCH_LIBS}
-        "${CMAKE_CURRENT_BINARY_DIR}/libtorchvision.so.1")
+    if(${RHEL_BUILD})
+      set(TRITON_PYTORCH_LIBS
+          ${TRITON_PYTORCH_LIBS}
+          "${CMAKE_CURRENT_BINARY_DIR}/libtorchvision.so")
+    else()
+      set(TRITON_PYTORCH_LIBS
+          ${TRITON_PYTORCH_LIBS}
+          "${CMAKE_CURRENT_BINARY_DIR}/libtorchvision.so.1")
+    endif()
   endif() # TRITON_PYTORCH_ENABLE_TORCHVISION
 
   if (${TRITON_PYTORCH_ENABLE_TORCHTRT})