/pytorch-cmake-example

Example to build PyTorch CUDA extension using CMake (with pybind11 and scikit-build)

Primary LanguageCuda

PyTorch CMake Example

This repository is an example for creating CMake-based pytorch CUDA extension. It's modified from pytorch extension example and scikit-build example.

To build this repository, install essential requirements and then execute python setup.py build. If you want to use CUDA in custom location (for example your library is installed from conda install cudatoolkit-dev -c conda-forge), you can give hint to CMake by defining CMake definition CMAKE_CUDA_COMPILER.