Add PyTorch lib directory to library path (#72)

openmm · Mar 7, 2022 · 84f7d88 · 84f7d88
1 parent e6a3179
commit 84f7d88
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/python/setup.py b/python/setup.py
@@ -7,21 +7,22 @@
 openmm_dir = '@OPENMM_DIR@'
 nn_plugin_header_dir = '@NN_PLUGIN_HEADER_DIR@'
 nn_plugin_library_dir = '@NN_PLUGIN_LIBRARY_DIR@'
+torch_dir, _ = os.path.split('@TORCH_LIBRARY@')
 
 # setup extra compile and link arguments on Mac
 extra_compile_args = ['-std=c++11']
 extra_link_args = []
 
 if platform.system() == 'Darwin':
     extra_compile_args += ['-stdlib=libc++', '-mmacosx-version-min=10.7']
-    extra_link_args += ['-stdlib=libc++', '-mmacosx-version-min=10.7', '-Wl', '-rpath', openmm_dir+'/lib']
+    extra_link_args += ['-stdlib=libc++', '-mmacosx-version-min=10.7', '-Wl', '-rpath', openmm_dir+'/lib', '-rpath', torch_dir]
 
 extension = Extension(name='_openmmtorch',
                       sources=['TorchPluginWrapper.cpp'],
                       libraries=['OpenMM', 'OpenMMTorch'],
                       include_dirs=[os.path.join(openmm_dir, 'include'), nn_plugin_header_dir],
                       library_dirs=[os.path.join(openmm_dir, 'lib'), nn_plugin_library_dir],
-                      runtime_library_dirs=[os.path.join(openmm_dir, 'lib')],
+                      runtime_library_dirs=[os.path.join(openmm_dir, 'lib'), torch_dir],
                       extra_compile_args=extra_compile_args,
                       extra_link_args=extra_link_args
                      )