I am thinking to define a new function, say, called F whose forward and backward methods can only be solved on CPUs.
So basically my F(x) requires reading out x.data.cpu().numpy(), and run some (sophiscated) code on CPUs, say g(x.data.cpu().numpy()). The backward pass for F also works similarly (and requires some saved numpy variables from the forward pass).
I am wondering whether I can implement similarly to other cases where the forward passes and backward passes use existing functions that are computed on GPUs?
I am tempted to just do the following in the forward pass
y = g(x.data.cpu().numpy())
Does that sound right? Thanks!