class MyReLU(nn.Module):
    def __init__(self):
        super(MyReLU, self).__init__()
        self.layer1 = nn.ReLU()
    
    def forward(self, x):
        return self.layer1(x)
    
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1 = nn.Sequential(nn.Linear(2,2), MyReLU())
        self.layer2 = nn.Sequential(nn.Linear(2,2), nn.ReLU())

    def forward(self, x):
        return self.layer2(self.layer1(x))
    
net = Net()
t = Variable(torch.ones(2), requires_grad=True)
traced, _ = torch.jit.trace(net, (t, ))
graph = torch._C._jit_get_graph(traced)

for n in graph.nodes():
    print('%s: %s' % (n.kind(), n.scopeName()))

t: Net/Sequential[layer1]/Linear[0]
unsqueeze: Net/Sequential[layer1]/Linear[0]
mm: Net/Sequential[layer1]/Linear[0]
squeeze: Net/Sequential[layer1]/Linear[0]
matmul: Net/Sequential[layer1]/Linear[0]
add: Net/Sequential[layer1]/Linear[0]
threshold: Net/Sequential[layer1]/MyReLU[1]/ReLU[layer1]
t: Net/Sequential[layer2]/Linear[0]
unsqueeze: Net/Sequential[layer2]/Linear[0]
mm: Net/Sequential[layer2]/Linear[0]
squeeze: Net/Sequential[layer2]/Linear[0]
matmul: Net/Sequential[layer2]/Linear[0]
add: Net/Sequential[layer2]/Linear[0]
threshold: Net/Sequential[layer2]/ReLU[1]