torch.autograd.detect_anomaly()

import torch

# 正向傳播時：開啟自動求導的異常偵測
torch.autograd.set_detect_anomaly(True)

# 反向傳播時：在求導時開啟偵測
with torch.autograd.detect_anomaly():
    loss.backward()