Solved by ptrblck in post #4 
              
              
                Your code works fine: 
class BranchRoutingModule(nn.Module):
  def __init__(self, in_channels=512, epsilon=1e-12):
    super(BranchRoutingModule, self).__init__()

    self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=1)
    self.l2norm = lambda x: torch.nn.functi…