Solved by ndronen in post #2 
              
              
                Here’s a little hack that does the trick: for each path the output of M takes, apply the identity function and call retain_grad on its output. 
#!/usr/bin/env python

# coding: utf-8

from functools import partial

import torch
import torch.nn as nn


def print_tensor_grad(grad, name=None, value=Non…