I guess your input to torch.atan2 might create invalid gradients as seen here:
x = torch.zeros(4, requires_grad=True)
y = torch.zeros(4)
out = torch.atan2(x, y)
out.mean().backward()
print(x.grad)
> tensor([nan, nan, nan, nan])

and as is pointed to in the error message.