def _check_nan_loss(self, loss):
    if torch.isnan(loss):
        norms = self.grad_norm(1)
        import json
        print(json.dumps(norms, indent=2))
def _check_nan_loss(self, loss):
    if torch.isnan(loss):
        norms = self.grad_norm(1)
        import json
        print(json.dumps(norms, indent=2))