Template for CUDA / C++ extension writing with PyTorch
- Python 3.6 (or more recent)
- pip
You can install the package using the default method:
pip install -e . --upgradewhich will build the CUDA extension and bind it with your python module.
This python package can be used like any other:
import torch
import cuda_ext.nn as nn
my_module = nn.DSigmoid()
with torch.no_grad():
out = my_module(torch.rand(2, 3, 32, 32).cuda())Please refer to CONTRIBUTING if you wish to contribute to this project.
The content of this repo was brought together by the repo owner but highly benefited from the following resources:
- CUDA extensions of Thomas Brandon
- Official PyTorch documentation
Distributed under the MIT License. See LICENSE for more information.