Try to use grid.device instead.



 sunshineatnoon:

But I think it’s inefficient to pass the grid every forward.


Might be and you should stick to your work flow, as I was just using it as an example. 
You could also try to register self.grid as a buffer using self.register_buffer, which would move the tensor automatically using model.to().