/torch_customop

An example of CUDA custom operation for PyTorch (weighted average)

Primary LanguageCuda

torch_customop

Just run the file by python main.py