Thanks for the code!


autograd.detect_anomaly() should be quite expensive, so if your code runs fine, I would recommend to remove this call.

Variable(batch.type(Tensor)): Variables are deprecated since 0.4.0, so you can just use tensors in newer versions. Also, if you want to convert the dtype, you should use to() or directly call the type as e.g. float().
I haven’t debugged your code, but do you really need to retain the graph in grad_x? I’m a bit worried it might be stored unnecessarily.
Since you are running the code on a GPU and CUDA calls are asynchronous, you should call torch.cuda.synchronize() before starting and stopping the timer. Otherwise your profiling might yield wrong results (e.g. kernel launch times)