The shape mismatch is raised in these layers:
            nn.Conv2d(128, 256, kernel_size = kernel_size, stride = stride, padding = padding),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size = kernel_size, stride = stride, padding = padding),

as the former conv layer creates an output activation with 256 channels while the latter expects 128 input channels.