I dont quite agree with you that we need any intermediate tensors created from
feature_target = model(input1)
feature_identity = model(input2)

The only tensors required for the loss computation are the feature_target and feature_identity, not their intermediate tensors.
I will use retain graph for now but it doesn’t feel right to me. Hopefully someone from the dev team can help to explain the best way to do this.