Solved by fmassa in post #2 
              
              
                That is one of the cases where it’s better not to use nn.Sequential, and inherit from model yourself and perform the operations that you want. 
For example 
class MyModel(nn.Module):
    def forward(self, input):
        return input ** 2 + 1
model = MyModel()

But if you want an equivalent to a Lam…