You need to place the input and target tensors on the same device as the model parameters.
There are advanced use cases for model sharding, which doesn’t seem to be used in your case.
If your validation loop increases the memory usage, you might want to either delete all unnecessary training tensors manually or wrap the training code in a train() function (and val() function), so that the temporary variables will be deleted automatically. Python uses function scoping so temp. variables, such as the model output, loss etc. will be deleted once you are leaving the function scope.