Hi asberman,
As I understand LayerNorm will compute mean and variance elementwise (not per batch), thus you should pass the spatial dimension of the input, not the channel dimension as in the case of BatchNorm.
Actually, I am doing the same work, and you can try to change the following:
the first layer norm :
nn.LayerNorm(num_disc_filters * 2), --> nn.LayerNorm([num_disc_filters * 2, 16, 16]),
the second:
nn.LayerNorm(num_disc_filters * 4), --> nn.LayerNorm([num_disc_filters * 4, 8, 8]),
the third:
nn.LayerNorm(num_disc_filters * 8), --> nn.LayerNorm([num_disc_filters * 8, 4, 4]),
I think this will work without error.
However, I found that even after replacing the BatchNorm with the LayerNorm, the model still diverges!!!
If you have solved this properly, please inform me.
Thanks