Fixes for OS X

.gitignore

@@ -4,6 +4,7 @@ torch.egg-info/
 */**/__pycache__
 torch/csrc/generic/TensorMethods.cpp
 torch/lib/*.so*
+torch/lib/*.dylib*
 torch/lib/*.h
 torch/lib/build
 torch/lib/tmp_install
@@ -12,3 +13,5 @@ torch/csrc/nn/THNN.cpp
 torch/csrc/nn/THCUNN.cwrap
 torch/csrc/nn/THCUNN.cpp
 */**/*.pyc
+*/**/*.so*
+*/**/*.dylib*

setup.py

@@ -151,9 +151,16 @@ def run(self):
 
 if WITH_CUDA:
     if platform.system() == 'Darwin':
-        include_dirs += ['/Developer/NVIDIA/CUDA-7.5/include']
+        cuda_path = '/Developer/NVIDIA/CUDA-7.5'
+        cuda_include_path = cuda_path + '/include'
+        cuda_lib_path = cuda_path + '/lib'
     else:
-        include_dirs += ['/usr/local/cuda/include']
+        cuda_path = '/usr/local/cuda'
+        cuda_include_path = cuda_path + '/include'
+        cuda_lib_path = cuda_path + '/lib64'
+    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']
     main_libraries += ['THC']
     main_sources += [
@@ -167,6 +174,13 @@ def run(self):
     extra_compile_args += ['-O0', '-g']
     extra_link_args += ['-O0', '-g']
 
+
+def make_relative_rpath(path):
+    if platform.system() == 'Darwin':
+        return '-Wl,-rpath,@loader_path/' + path
+    else:
+        return '-Wl,-rpath,$ORIGIN/' + path
+
 ################################################################################
 # Declare extensions and package
 ################################################################################
@@ -180,7 +194,7 @@ def run(self):
     language='c++',
     extra_compile_args=extra_compile_args,
     include_dirs=include_dirs,
-    extra_link_args=extra_link_args + ['-Wl,-rpath,$ORIGIN/lib'],
+    extra_link_args=extra_link_args + [make_relative_rpath('lib')]
 )
 extensions.append(C)
 
@@ -190,7 +204,7 @@ def run(self):
     language='c++',
     extra_compile_args=extra_compile_args,
     include_dirs=include_dirs,
-    extra_link_args=extra_link_args + ['-Wl,-rpath,$ORIGIN/../lib'],
+    extra_link_args=extra_link_args + [make_relative_rpath('../lib')]
 )
 extensions.append(THNN)
 
@@ -201,7 +215,7 @@ def run(self):
         language='c++',
         extra_compile_args=extra_compile_args,
         include_dirs=include_dirs,
-        extra_link_args=extra_link_args + ['-Wl,-rpath,$ORIGIN/../lib'],
+        extra_link_args=extra_link_args + [make_relative_rpath('../lib')]
     )
     extensions.append(THCUNN)
     packages += ['torch.cuda', 'torch.legacy.cunn']

torch/csrc/cuda/Module.cpp

@@ -177,7 +177,7 @@ PyObject * THCPModule_initExtension(PyObject *self)
   PyObject *torch_module = PyImport_ImportModule("torch.cuda");
   if (!torch_module) {
     THPUtils_setError("class loader couldn't access torch module");
-    return false;
+    return NULL;
   }
   PyObject* module_dict = PyModule_GetDict(torch_module);
   return PyBool_FromLong(THCPModule_initCuda(module_dict));

torch/lib/build_all.sh

@@ -4,8 +4,14 @@ set -e
 BASE_DIR=$(pwd)
 cd torch/lib
 INSTALL_DIR=$(pwd)/tmp_install
-BASIC_FLAGS=" -DTH_INDEX_BASE=0 -I$INSTALL_DIR/include -I$INSTALL_DIR/include/TH -I$INSTALL_DIR/include/THC -L$INSTALL_DIR/lib "
-FLAGS="$BASIC_FLAGS -Wl,-rpath,\$ORIGIN"
+BASIC_FLAGS=" -DTH_INDEX_BASE=0 -I$INSTALL_DIR/include -I$INSTALL_DIR/include/TH -I$INSTALL_DIR/include/THC "
+LDFLAGS="-L$INSTALL_DIR/lib "
+if [[ $(uname) == 'Darwin' ]]; then
+    LDFLAGS="$LDFLAGS -Wl,-rpath,@loader_path"
+else
+    LDFLAGS="$LDFLAGS -Wl,-rpath,\$ORIGIN"
+fi
+FLAGS="$BASIC_FLAGS $LDFLAGS"
 function build() {
   mkdir -p build/$1
   cd build/$1
@@ -18,6 +24,15 @@ function build() {
               -DTH_INCLUDE_PATH="$INSTALL_DIR/include"
   make install -j$(getconf _NPROCESSORS_ONLN)
   cd ../..
+
+  if [[ $(uname) == 'Darwin' ]]; then
+    cd tmp_install/lib
+    for lib in *.dylib; do
+      echo "Updating install_name for $lib"
+      install_name_tool -id @rpath/$lib $lib
+    done
+    cd ../..
+  fi
 }
 
 mkdir -p tmp_install