That’s basically right. If you use loss after the condition to call loss.backward() on it, it will create an error though.
I would rather zero out the loss:
loss = calculate the loss
loss = loss * 0.
Alternatively your training code would have to be in the if STATEMENT condition.