Hi,
You need to wrap your forward and backward pass with it, so wherever you do that.
If you want it for the whole script, you can also add torch.autograd.set_detect_anomaly(True) at the beginning of your script and it will stay on !