You have to observe the torch.autograd.Function protocol, so forward operates on tensors and backward on Variables (until you use a version where they are merged). There are good examples floating around, if you are looking for one, I can offer a differentiable Implicit Function 10 implementation.
I would advise to either require the caller to move it to CPU (and just raise NotImplemented if they are on GPU - that would be my preferred method given that moving to CPU has a fairly heavy performance impact as it stops the asynchronous GPU processing pytorch does) or move it back to where the caller had them.
Best regards
Thomas