@ptrblck this should be the model without permute.
class SimpleBinNetNoPermute(nn.Module):
    def __init__(self):
        super(SimpleBinNetNoPermute, self).__init__()

        conv1 = self.make_conv2d_layer(1, 8, (1, 9), (1, 2))
        conv2 = self.make_binconv2d_layer(8, 64, (1, 5), (1, 2))
        conv3 = self.make_binconv2d_layer(64, 32, (1,3), padding=1)

        sfebs = conv1 + conv2 + [nn.MaxPool2d(kernel_size=(1, 50))]
        self.sfeb = nn.Sequential(
            *sfebs
        )

        tfebs = conv3 + [nn.MaxPool2d(kernel_size = (1,2))]
        self.tfeb = nn.Sequential(
            *tfebs
        )

        fcn = nn.Linear(7200, 10, bias=True)
        fcn.apply(lambda m: nn.init.kaiming_normal_(m.weight, nonlinearity='sigmoid'))
        self.output = nn.Sequential(
            nn.Flatten(),
            fcn,
            nn.Softmax(dim=1)
        )

        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                if hasattr(m.weight, 'data'):
                    m.weight.data.zero_().add_(1.0)

    def forward(self, x):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                if hasattr(m.weight, 'data'):
                    m.weight.data.clamp_(min=0.01)
        x = self.sfeb(x)
        x = self.tfeb(x)
        print(x.shape)
        y = self.output(x)
        return y

    def make_conv2d_layer(self, in_ch, out_ch, k_size, stride=(1,1), padding=0):
        conv = nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=k_size, stride=stride, padding=padding, bias=False)
        conv.apply(lambda m: nn.init.kaiming_normal_(m.weight, nonlinearity='relu'))
        return [
            conv,
            nn.BatchNorm2d(out_ch),
            nn.ReLU()
        ]
    def make_binconv2d_layer(self, in_ch, out_ch, k_size, stride=(1,1), padding=0):
        return [
            nn.BatchNorm2d(in_ch),
            BinConv2d(ch_in=in_ch, ch_out=out_ch, kernel_size=k_size, stride=stride, padding=padding),
            nn.ReLU()
        ]