Running custom CUDA code using PyTorch

I am using PyCUDA, and recently started using PyTorch due to its native GPU tensor support.

I would like to write a new PyTorch method, lets say converting an RGB image to BW (see the PyCUDA equivalent here: that uses CUDA directly.

Is that possible?


1 Like