A better way I’ve found to change the graph after the fact is to:

.detach() the model outputs,
use those to alter the targets(these aren’t in the graph and can be altered at this stage before the loss function),
Run your loss function on the altered targets.

You won’t need to take the outputs or targets off GPU. Just keep all of the tensors on GPU and run your calcs there for speed.