For improved Wasserstein GAN (aka Wasserstein GAN with gradient penalty [WGAN-GP]), layer normalization is recommended in the discriminator, as opposed to nn.BatchNorm2d
I see that nn.LayerNorm was (relatively) recently added to torch.nn.modules, and I’d like to use it, as opposed to writing my own layer normalization. However, just replacing calls to nn.BatchNorm2d(input_size)  with nn.LayerNorm(input_size) gives the following error:
File "/usr/local/lib/python3.6/site-packages/torch/nn/functional.py", line 1314, in layer_norm
PSTtorch.backends.cudnn.enabled)
PSTRuntimeError: Given normalized_shape=[128], expected input with shape [*, 128], 
but got input of size[128, 128, 16, 16]

This is the structure of the network:
class WgangpDiscriminator(nn.Module):  
    def __init__(self, channels, num_disc_filters):
        super(WgangpDiscriminator, self).__init__()
        self.ngpu = 1
        self.main = nn.Sequential(
        #num_disc_filters = 64
        nn.Conv2d(channels, num_disc_filters, 4, 2, 1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),

        nn.Conv2d(num_disc_filters, num_disc_filters * 2, 4, 2, 1, bias=False),
        nn.LayerNorm(num_disc_filters * 2),
        nn.LeakyReLU(0.2, inplace=True),

        nn.Conv2d(num_disc_filters * 2, num_disc_filters * 4, 4, 2, 1, bias=False),
        nn.LayerNorm(num_disc_filters * 4),
        nn.LeakyReLU(0.2, inplace=True),

        # state size. (num_disc_filters*4) x 8 x 8
        # nn.Conv2d(num_disc_filters * 4, num_disc_filters * 8, 4, 2, 1, bias=False),
        # nn.BatchNorm2d(num_disc_filters * 8),
        # nn.LeakyReLU(0.2, inplace=True),

        # was num_disc_filters * 16
        nn.Conv2d(num_disc_filters * 4, num_disc_filters * 8, 4, 2, 1, bias=False),
        nn.LayerNorm(num_disc_filters * 8),
        nn.LeakyReLU(0.2, inplace=True),
        
        # was num_disc_filters * 16
        nn.Conv2d(num_disc_filters * 8, 1, 4, 1, 0, bias=False),
        # ty  https://github.com/pytorch/examples/issues/70 apaske
    )

    def forward(self, inp):
         if isinstance(inp.data, torch.cuda.FloatTensor) and self.ngpu > 1:
             output = nn.parallel.data_parallel(self.main, inp, range(self.ngpu))
         else:
             output = self.main(inp)
         return output.view(-1, 1).squeeze(1)

nn.LayerNorm expects normalized_shape as input ( an int, list or torch.Size), but nn.Conv2d layers don’t have .size , .get_shape() or .shape(), so I can’t follow the example in the docs:
input = torch.randn(20, 5, 10, 10) 
# With Learnable Parameters 
m = nn.LayerNorm(input.size()[1:])

How do I do this conversion?
Many thanks in advance!