Custom Autograd.Functions to GPU

I am just starting out, and I need to create a custom activation function just like how they did it here:

So my question is does relu(input) run on the cpu or gpu? I am guessing that it doesn’t run on the gpu, since anytime I want something to run on the gpu, I always had to call .to(device).

I tried that with the custom activation, but I was unable to do that. So if I am correct in that the activation does not run on the gpu, then how do I do that.

If I am wrong, can you please explain how it knows to send to gpu?


These Functions don’t have any parameters, so they will work with whatever inputs are given. Just give it inputs on the gpu and it will run on the gpu.