Take the former example
layer1 = nn.Conv2d(param1)
layer2 = nn.Conv2d(param2)
net = nn.Sequential([layer1, layer2])

output = net(input)

layer_grad_fn = output.grad_fn.next_functions[0][0]

Here In the graph I need to know the number of input and output channels of layer_grad_fn, which is a ThnnConv2DBackward. And sometimes the conv kernel size as well.
Is there any solution?