If you cannot or don’t want to register these tensors as parameters of buffers, you could manually move them to the corresponding device by using the .device attribute of the input tensor:
def forward(self, x):
my_tensor = my_tensor.to(x.device)