After having a brief look through, it seems you’re swapping between torch and numpy, when moving back and forth between the library would break the gradient of any intermediate computations, no? I’d suggest trying to remove all dependencies on numpy  and purely use torch operations so autograd can track the operations.
Also, you do use the gradient of your input data at all (i.e. x) ? If you do, make sure to enable grad for that data! You can add x.requires_grad_() before your loop.