Hi,
The error seems to be that you tried to call backward twice on the same graph without specifying retain_variables=True during the first backward.
You can also try to change dist = dist[dist > 0.] with dist = dist[(dist > 0.).detatch()] to force the indices to be properly detached.
In any case, your code should work so please let us know if explicitly detaching solves the issue.