From 3828b5235e68e6035d414cd1a0e15e85be288586 Mon Sep 17 00:00:00 2001 From: Peter Eastman Date: Wed, 29 Jan 2025 15:50:59 -0800 Subject: [PATCH] Load correct library name --- CMakeLists.txt | 4 ++-- src/pytorch/__init__.py | 9 ++++++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 89f34917..a68df575 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -32,12 +32,12 @@ set(SRC_FILES src/ani/CpuANISymmetryFunctions.cpp # Build the library set(LIBRARY ${NAME}PyTorch) add_library(${LIBRARY} SHARED ${SRC_FILES}) -set_property(TARGET ${LIBRARY} PROPERTY CXX_STANDARD 14) +set_property(TARGET ${LIBRARY} PROPERTY CXX_STANDARD 17) target_include_directories(${LIBRARY} PRIVATE ${Python3_INCLUDE_DIRS} src/ani src/pytorch src/schnet) target_link_libraries(${LIBRARY} ${TORCH_LIBRARIES} ${Python3_LIBRARIES}) if(ENABLE_CUDA) - set_property(TARGET ${LIBRARY} PROPERTY CUDA_STANDARD 14) + set_property(TARGET ${LIBRARY} PROPERTY CUDA_STANDARD 17) target_compile_definitions(${LIBRARY} PRIVATE ENABLE_CUDA) endif(ENABLE_CUDA) diff --git a/src/pytorch/__init__.py b/src/pytorch/__init__.py index 645d4399..b5ca42f8 100644 --- a/src/pytorch/__init__.py +++ b/src/pytorch/__init__.py @@ -2,9 +2,16 @@ High-performance PyTorch operations for neural network potentials ''' import os.path +import platform import torch -torch.ops.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so')) +if platform.system() == 'Darwin': + libname = 'libNNPOpsPyTorch.dylib' +elif platform.system() == 'Windows': + libname = 'NNPOpsPyTorch.dll' +else: + libname = 'libNNPOpsPyTorch.so' +torch.ops.load_library(os.path.join(os.path.dirname(__file__), libname)) from NNPOps.OptimizedTorchANI import OptimizedTorchANI