Solved by tom in post #2 
              
              
                The easiest probably is to just ask grad / backward for the parameter it got (after fixing the retain_grad vs. retain_graph typo): 
from torch import nn, is_grad_enabled, rand, randint


def _hook(module, grad_input, grad_output):
    fr = ([f[0] for f in inspect.stack() 
           if '/torch/autog…