Solved by ptrblck in post #6 
              
              
                If the batch norm layers are between linear layers, the shape should most likely once be changed before the first linear layer. 
Anyway, you could define a Permute layer and use it inside your nn.Sequential container: 
class Permute(nn.Module):
    def __init__(self, dims):
        super(Permute, se…