Solved by gbilla in post #2 
              
              
                Fixed it. Key takeaways in fixing autograd problems like these: 

Use torch.autograd.set_detect_anomaly(True) to isolate the line with the problem.
Call .clone() on the variables on the right hand side of the update.

had to replace the line: 
dx = 2.0 * torch.stack([torch.rfft(x_r[0], 2), torch.rff…