@colesbury Thanks for your help. I tried your suggestion and when I print output.grad it is only consisting of zeros and nans. Per your comment, I removed the clamp operation from my loss function and looked at the grad again and it looks like this time I see some numbers and some nans. More importantly, this time the optimizer is able to reduce my loss function. So I think your intuition is correct and clamp messes up everything.
Thanks again for your help.