It depends on what you are trying to achieve and I wouldn’t claim the approach is fundamentally flawed, as I’m not familiar with the use case of @adm.



 raraz15:

This makes it not possible to use BatchNorm and Linear layer in the sequential format. At least to my knowledge.


You can use nn.BatchNorm1d layers in an nn.Sequential container as seen here:
model = nn.Sequential(
    nn.Linear(4, 3),
    nn.BatchNorm1d(3)
)

x = torch.randn(2, 4)
out = model(x)

If you add another dimension and want to permute the activation, you could write a custom module, which permutes it to the right shape.