Assume that I have a model that has two parameters A and B, takes as input a matrix C, and returns f(g(A, B), C). I could compute this during every forward fun, but we really should be able to cache g(A, B). I tried to compute self.D = g(A, B) in __init__, but the device doesn’t match when I do f(self.D, C). Alternatively, if I do f(self.D.to(C.device), C), the forward pass works but the gradient device is wrong during the optimizer step. I think I could also somehow always keep self.D in GPU memory at all times, but is that necessary? What’s the best way to accomplish this?