Set LD_LIBRARY_PATH in install script (#2671)

Summary:
- Set LD_LIBRARY_PATH in the fbgemm_gpu install script to fix the issue with fbgem_gpu being unable to locate libnvrtc.so on start

Pull Request resolved: #2671

Reviewed By: q10, iamzainhuda

Differential Revision: D67949409

Pulled By: TroyGarden

fbshipit-source-id: dabef28755ee123d9560d99848b2764598e9b3a6

.github/scripts/install_fbgemm.sh

@@ -5,11 +5,17 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-echo "CU_VERSION"
-echo "$CU_VERSION"
+echo "CU_VERSION: ${CU_VERSION}"
+echo "CHANNEL: ${CHANNEL}"
+echo "CONDA_ENV: ${CONDA_ENV}"
 
-echo "CHANNEL"
-echo "$CHANNEL"
+if [[ $CU_VERSION = cu* ]]; then
+    # Setting LD_LIBRARY_PATH fixes the runtime error with fbgemm_gpu not
+    # being able to locate libnvrtc.so
+    echo "[NOVA] Setting LD_LIBRARY_PATH ..."
+    conda env config vars set -p ${CONDA_ENV}  \
+        LD_LIBRARY_PATH="/usr/local/lib:${CUDA_HOME}/lib64:${CONDA_ENV}/lib:${LD_LIBRARY_PATH}"
+fi
 
 if [ "$CHANNEL" = "nightly" ]; then
     ${CONDA_RUN} pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/"$CU_VERSION"