Fix build issue. CUDA may be installed in $CUDA_HOME/lib on macOS. …

…(#23491)

Summary:
Closes gh-16955.
Closes pytorch/vision#977

On Linux both `lib64` and `lib` may be present (symlinked). The reports
seem to all be about macOS, but it seems like this is also possibly more
robust on Linux and can't hurt. So not treating platforms differently.

Note that Eigen has a similar check in its CMake:

```
if(CUDA_64_BIT_DEVICE_CODE AND (EXISTS "${CUDA_TOOLKIT_ROOT_DIR}/lib64"))
  link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
else()
  link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib")
endif()
 ```

There may be other issues for building from source on macOS, can't test.
Pull Request resolved: pytorch/pytorch#23491

Differential Revision: D16538973

Pulled By: soumith

fbshipit-source-id: cc309347b7d16e718e06878d3824d0a6e40b1019

torch/utils/cpp_extension.py

@@ -548,7 +548,17 @@ def library_paths(cuda=False):
         paths.append(lib_path)
 
     if cuda:
-        lib_dir = 'lib/x64' if IS_WINDOWS else 'lib64'
+        if IS_WINDOWS:
+            lib_dir = 'lib/x64'
+        else:
+            lib_dir = 'lib64'
+            if (not os.path.exists(_join_cuda_home(lib_dir)) and
+                    os.path.exists(_join_cuda_home('lib'))):
+                # 64-bit CUDA may be installed in 'lib' (see e.g. gh-16955)
+                # Note that it's also possible both don't exist (see
+                # _find_cuda_home) - in that case we stay with 'lib64'.
+                lib_dir = 'lib'
+
         paths.append(_join_cuda_home(lib_dir))
         if CUDNN_HOME is not None:
             paths.append(os.path.join(CUDNN_HOME, lib_dir))