Hi,
I need to access to modules of model when I take a backward graph trace. I want to access to graph trace and nn.Module simultanious. There are any ways to match between .grad_fn and nn.Module?
For example:


Import pytorch
import torch
from torch.autograd import Variable
from torch import nn


Create model (It can be more difficult model)
model = nn.Conv2d(3, 10, kernel_size=1)


Compute model for some input variable
input_var = Variable(torch.randn(1, 3, 10, 10), requires_grad=True)
output_var = model(input_var)


So, I can access to output_var.grad_fn now, it is a ConvNdBackward object. And I want to cast ConvNdBackward to nn.Conv2d and get access to all members of nn.Conv2d, and I want to have access to each submodule of model. Can I do it?
print(output_var.grad_fn)


I want to do something like this:
fn = output_var.grad_fn
conv = fn.to_module() # (grad_fn.module() -> nn.Module)