Solved by ptrblck in post #2 
              
              
                Most likely self.f is not pushed to the right device, since it is neither registered as an nn.Parameter nor as a buffer (using self.register_buffer). 
Use the former case, if self.f should require gradients and the latter if not. 
The manual cuda() call inside your __init__ method won’t work, as the…