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, self).__init__()
        self.dims = dims
        
    def forward(self, x):
        x = x.permute(self.dims).contiguous()
        return x