diff --git a/tensorflow/tools/ci_build/builds/test_user_ops.sh b/tensorflow/tools/ci_build/builds/test_user_ops.sh
index 4751157efee8b6fa656686b6f8b548a3c645f132..216abbe8e6797b8ab1107e2cdee4741bfa0c6108 100755
--- a/tensorflow/tools/ci_build/builds/test_user_ops.sh
+++ b/tensorflow/tools/ci_build/builds/test_user_ops.sh
@@ -184,11 +184,23 @@ else
       -I "${TF_INC}" -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC || \
       die "nvcc compilation of ${OP_KERNEL_CC} FAILED"
 
+  CUDA_LIB_DIR="/usr/local/cuda/lib64"
+  if [[ ! -d "${CUDA_LIB_DIR}" ]]; then
+    CUDA_LIB_DIR="/usr/local/cuda/lib"
+  fi
+  if [[ ! -d "${CUDA_LIB_DIR}" ]]; then
+    die "ERROR: Failed to find CUDA library directory at either of "
+"/usr/local/cuda/lib and /usr/local/cuda/lib64"
+  fi
+
+  echo "Found CUDA library diretory at: ${CUDA_LIB_DIR}"
+  echo ""
+
   # USER_OP_SO=$(basename $(echo "${OP_KERNEL_CC}" | sed -e 's/\.cc/\.so/'))
   USER_OP_SO="add_one.so"
   "${GPP_BIN}" -std=c++11 ${EXTRA_GPP_FLAGS} \
       -shared -o "${USER_OP_SO}" "${OP_KERNEL_CC}" \
-      "${OP_KERNEL_O}" -I "${TF_INC}" -L "/usr/local/cuda/lib64" \
+      "${OP_KERNEL_O}" -I "${TF_INC}" -L "${CUDA_LIB_DIR}" \
       -fPIC -lcudart || \
       die "g++ compilation of ${OP_KERNEL_CC}" FAILED
 fi