well, yes, I am of course not sure, what I am doing. But I mean, I can use any kind of loss in tensorflow, so it is something special with autograd. If it was me, I would simply have output = net(input), backpropagate within there and it’s fine.
My second approach is very simple. Copying the output and putting the input tensor values inside, while dropping the unused values. Then using the modified output in the loss. So kinda hacking the output to match the input values, while keeping the graph. Gradients are changing at least, going to inf and nan quickly though with SGD