The number of features for an nn.BatchNorm2d layer should be equal to the number of output channels of the preceding conv layer.
Fix the model definition from e.g.
[nn.Conv2d(128, 256, 5, stride=2, padding=2, bias=False), norm(128), act]

to
[nn.Conv2d(128, 256, 5, stride=2, padding=2, bias=False), norm(256), act]

PS: I’ve formatted your code for better readability. If you would like to post code snippets, you can wrap your code in three backticks ```