I’m not sure what your actual use case is, but you could directly use torch methods to write a new layer.
As long as you are using differentiable operations, Autograd will capture the computation graph and will be able to compute the backward pass.
If you are breaking the computation graph for some reason, you would need to implement the backward method manually as described here.
Are you writing a custom CUDA kernel or should the layer just be executable on the GPU?