Try to isolate the iteration which causes this issue and check the inputs as well as outputs to torch.pow. Based on your code I cannot find anything obviously wrong. Also, I would recommend to post code snippets directly by wrapping them into three backticks ``` (as you’ve already done), as it would make debugging easier and the search engine could also use it for results in case other users face similar issues.