Output of the module, computed with the last call

Is there an equivalent to torch module output member in pytorch?

For instance, looking at the following class

class Example(nn.Module):
    def __init__(self):
        self.L1 = nn.Linear(1,10)
        self.L2 = nn.Linear(10,1)
    def forward(self,x)
          x = self.L2(self.L1(x))
          return x

Is it possible that after I call forward by

m = Example()

to get access to m.L1 output for this run?


No, what you can do is

def forward(self, x):
    x = self.L1(x)
    y = self.L2(x)
    return x,y

By design, pytorch modules hold the parameters, but do not “contain the computation”.

Best regards


I can understand the logic, but I wonder if this is the best option. My intentions are for debugging and visualizing the computations (e.g. Attention units).
My model contains large, not constant, number of modules, so the suggested solution won’t be very elegant for my needs. Maybe a different solution exists, like a wrapper module that saves the computations or hacks?