detect_anomaly.py

import torch

def main():
    torch.autograd.detect_anomaly()
    ...
    # Rest of the training code
   

# OR
class MyNumericallyUnstableLoss(nn.Module):
  def forward(self, input, target):
    with torch.autograd.set_detect_anomaly(True):
       loss = input * target
       return loss