Hi,
You don’t actually need to clone here.  You can re-use z as the input for f multiple times with no issues.
Your weights do seem to change in both branches. Do you have any issue with the loss value?