shivangi:

But I am getting the following error:
kl_div = F.kl_div(p_temp, q_temp)
RuntimeError: the derivative for ‘target’ is not implemented


This is quite literally what it says: F.kl_div does not support taking gradients w.r.t. the second (target) argument and you q_temp asks for gradients (requires_grad=True). If you want the derivative for the second argument as well, you would need to spell out the kl_div, too.
If I may say so, you likely want to replace the for loops over coordinates with clever broadcasting 7, this is usually much, much faster to the point where it is an error to do for loops over coordinates (very narrow exceptions apply when you would have unreasonably large intermediates or somesuch).
Best regards
Thomas