@albanD Thanks for your interest
I have set your “magic” line right after the imports
import torch
import torch.nn as nn
from torch.autograd import grad
import numpy as np

torch.autograd.set_detect_anomaly(True) 

Here is the Traceback
Warning: Traceback of forward call that caused the error:
  File "...anaconda3/lib/python3.7/traceback.py", line 197, in format_stack
    return format_list(extract_stack(f, limit=limit))
 (print_stack at /opt/conda/conda-bld/pytorch_1579022060824/work/torch/csrc/autograd/python_anomaly_mode.cpp:57)
Traceback (most recent call last):
  File "bugs.py", line 335, in <module>
    loss.backward()
  File "...anaconda3/lib/python3.7/site-packages/torch/tensor.py", line 195, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "...anaconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Expected a Tensor of type Variable but found an undefined Tensor for argument #0 'self'