So new infos. I have run the code on Google Colab.
It uses Python 3.6, Torch 1.3.1 and crash both on CPU and GPU K80 with the same message.
/pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:57: UserWarning: Traceback of forward call that caused the error:
  File "/usr/lib/python3.6/traceback.py", line 197, in format_stack
    return format_list(extract_stack(f, limit=limit))

---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

<ipython-input-1-e5464bc91872> in <module>()
    333 print("loss   : ",loss)
    334 #compute grad
--> 335 loss.backward()
    336 
    337 print("end")

1 frames

/usr/local/lib/python3.6/dist-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     97     Variable._execution_engine.run_backward(
     98         tensors, grad_tensors, retain_graph, create_graph,
---> 99         allow_unreachable=True)  # allow_unreachable flag
    100 
    101 

RuntimeError: sizes() called on undefined Tensor

May be @albanD you have an idea?