Change:
right now if is not in LDLIBRARYPATH, though it was found at compile-time, these lines will fail: Avoid this, by idk doing something... I think i have a few good ideas.
We need to add /usr/local/cuda/lib64 (or that's mac dir) to the rpath of _C
the problem i dont think is to add it to rpath, it's because we
Right. We could try with the default cuda installation path if it's not found from LDLIBRARYPATH.
Actually, do we really want to use the compile path? We should try to detect it at runtime, right? Or are we going to ship CUDA libs in our binaries?
we first load using runtime paths, and if it fails, fallback to path known at compile time.
I fear this is going to get brittle, especially if you are shipping CUDA libs. e.g. a user swaps HW and picks up a new driver and/or toolkit. As says, I'm pretty sure you want to use the runtime paths. Similarly, CUDA 7.5 won't run on Pascal and we have a cudnn built for 7.5 that will not do the right thing on CUDA 8, so those need to go together. If a user installs the CUDA bits via deb/rpm the right things are supposed to happen. If you push bits, we need to figure out how to make sure things line up. But this is perhaps a much larger packaging conversation.
Yeah, I've implemented it to first try loading it without any path and then try and . I'll push the commit tomorrow.

Discussion:
import sys import os <del> # TODO: make this more robust WITH_CUDA = os.path.exists('/Developer/NVIDIA/CUDA-7.5/include') or os.path.exists('/usr/local/cuda/include') </del> <ins> CUDA_HOME = os.getenv('CUDA_HOME', '/usr/local/cuda') WITH_CUDA = os.path.exists(CUDA_HOME) </ins> WITH_CUDNN = WITH_CUDA DEBUG = False

pass if WITH_CUDA: <del> if platform.system() == 'Darwin': cuda_path = '/Developer/NVIDIA/CUDA-7.5' cuda_include_path = cuda_path + '/include' cuda_lib_path = cuda_path + '/lib' else: cuda_path = '/usr/local/cuda' cuda_include_path = cuda_path + '/include' cuda_lib_path = cuda_path + '/lib64' </del> <ins> cuda_lib_dirs = ['lib64', 'lib'] cuda_include_path = os.path.join(CUDA_HOME, 'include') for lib_dir in cuda_lib_dirs: cuda_lib_path = os.path.join(CUDA_HOME, lib_dir) if os.path.exists(cuda_lib_path): break </ins> include_dirs.append(cuda_include_path) extra_link_args.append('-L' + cuda_lib_path) extra_link_args.append('-Wl,-rpath,' + cuda_lib_path) extra_compile_args += ['-DWITH_CUDA'] <ins> extra_compile_args += ['-DCUDA_LIB_PATH=' + cuda_lib_path] </ins> main_libraries += ['THC'] main_sources += [ "torch/csrc/cuda/Module.cpp",

extern PyObject * THCPModule_initialSeed(PyObject *_unused); extern PyObject * THCPModule_cudaHostAllocator(PyObject *_unused); extern PyObject * THCPModule_cudaSynchronize(PyObject *_unused); <ins> extern PyObject * THCPModule_getLibPath(PyObject *_unused); </ins> #endif static PyMethodDef TorchMethods[] = {

{"_cuda_initialSeed", (PyCFunction)THCPModule_initialSeed, METH_NOARGS, NULL}, {"_cuda_cudaHostAllocator", (PyCFunction)THCPModule_cudaHostAllocator, METH_NOARGS, NULL}, {"_cuda_synchronize", (PyCFunction)THCPModule_cudaSynchronize, METH_NOARGS, NULL}, <ins> {"_cuda_getLibPath", (PyCFunction)THCPModule_getLibPath, METH_NOARGS, NULL}, </ins> #endif {"_safe_call",      (PyCFunction)THPModule_safeCall,          METH_VARARGS | METH_KEYWORDS, NULL}, {"_sendfd",         (PyCFunction)THPModule_sendfd,            METH_VARARGS, NULL},

END_HANDLE_TH_ERRORS } <ins> PyObject * THCPModule_getLibPath(PyObject *_unused) { #define _STR(x) #x #define STR(x) _STR(x) #if PY_MAJOR_VERSION == 2 return PyString_FromString(STR(CUDA_LIB_PATH)); #else return PyUnicode_FromString(STR(CUDA_LIB_PATH)); #endif #undef STR #undef _STR } </ins> //////////////////////////////////////////////////////////////////////////////// // Cuda module initialization ////////////////////////////////////////////////////////////////////////////////

import contextlib import platform import ctypes <ins> import os </ins> import torch _initialized = False

return (hasattr(torch._C, '_cuda_isDriverSufficient') and torch._C._cuda_isDriverSufficient()) <del> def _lazy_init(): global _initialized, _cudart if _initialized: return </del> <ins> def _load_cudart(): system = platform.system() lib_name = 'libcudart.' + ('dylib' if system == 'Darwin' else 'so') lib_paths = [ lib_name, os.path.join(torch._C._cuda_getLibPath(), lib_name), os.path.join('/usr/local/cuda/lib64', lib_name), os.path.join('/usr/local/cuda/lib', lib_name), ] for path in lib_paths: try: return ctypes.cdll.LoadLibrary(path) except OSError: pass raise RuntimeError("couldn't find libcudart. Make sure CUDA libraries " "are installed in a default location, or that they're in " + ("DYLD_LIBRARY_PATH" if system == 'Darwin' else "LD_LIBRARY_PATH") + ".") def _check_driver(): </ins> if not hasattr(torch._C, '_cuda_isDriverSufficient'): raise AssertionError("Torch not compiled with CUDA enabled") if not torch._C._cuda_isDriverSufficient():

Alternatively, go to: https://pytorch.org/binaries to install a PyTorch version that has been compiled with your version of the CUDA driver.""".format(str(torch._C._cuda_getDriverVersion()))) <ins> def _lazy_init(): global _initialized, _cudart if _initialized: return _check_driver() </ins> assert torch._C._cuda_init() <del> _initialized = True if platform.system() == 'Darwin': _cudart = ctypes.cdll.LoadLibrary('libcudart.dylib') else: _cudart = ctypes.cdll.LoadLibrary('libcudart.so') </del> <ins> _cudart = _load_cudart() </ins> _cudart.cudaGetErrorName.restype = ctypes.c_char_p _cudart.cudaGetErrorString.restype = ctypes.c_char_p <ins> _initialized = True </ins> def cudart():