Solved by SimonW in post #2 
              
              
                I’ve used this, and it works fine on Cifar10 
class AllConvNet(nn.Module):

    def __init__(self, dropout=True, nc=3, num_classes=10):
        super(AllConvNet, self).__init__()
        self.dropout = dropout
        self.conv1 = nn.Conv2d(nc, 96, 3, padding=1)
        self.conv2 = nn.Conv2d(96, 96…