libtorch_custom_function

Simple example of a custom function for libtorch with a custom cuda kernel

Pre-build patching

For the CUDA version of libtorch, a change needs to be made to the CMake files within the Caffe2 subdirectory at libtorch/share/cmake/Caffe2/public/cuda.cmake Change the command used to compute the short-hash of libnvtrc to the python executable.

# Line 135
COMMAND Python::Interpreter -c

# Changed to
COMMAND ${Python_EXECUTABLE} -c

Build

mkdir build && cd build
source ~/Desktop/load_cuda.sh
export CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda-12.9
cmake ..