I am still having trouble with this  I think I have an easier example… If you use just one linear(1,1) module everything works as expected, but if you use one conv2d on a 1,1 input (equivalent to the linear module) then you get no gradient on the weights in the convolution…
import torch                                                          
from torch import nn                                                  
from torchviz import make_dot, make_dot_from_trace                    
                                                                      
model = nn.Sequential()                                               

#if its only one conv module in the model
#model.add_module('C0', nn.Conv2d(1,1,1))  #either uncomment this line                   

#if its just one linear module in the model
model.add_module('W0', nn.Linear(1,1))     #or uncomment this line
                                                                      
x = torch.randn(1,1,1,1).requires_grad_(True)                         
                                                                      
def double_backprop(inputs, net):                                     
    y = net(x).mean()**2                              
    grad,  = torch.autograd.grad(y, x, create_graph=True, retain_graph=True)   
    return grad.mean()                                                
    
dot=make_dot(double_backprop(x, model), params=dict(list(model.named_parameters()) + [('x', x)]))

dot.render(view=True)