How to add bayesian conv layers( getting this error RuntimeError: Expected 3-dimensional input for 3-dimensional weight 1 28 3, but got 4-dimensional input of size [100, 1, 32, 32] instead)

def main():
# import data
kwargs = {‘num_workers’: 1, ‘pin_memory’: True} if FLAGS.cuda else {}

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),lambda x: 2 * (x - 0.5),
                   ])),
    batch_size=FLAGS.batchsize, shuffle=True, **kwargs)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=False, transform=transforms.Compose([
        transforms.ToTensor(), lambda x: 2 * (x - 0.5),
    ])),
    batch_size=FLAGS.batchsize, shuffle=True, **kwargs)

# for later analysis we take some sample digits
mask = 255. * (np.ones((1, 28, 28)))
examples = train_loader.sampler.data_source.train_data[0:5].numpy()
images = np.vstack([mask, examples])

# build a simple MLP
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # activation
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        # layers
        self.conv1 = BayesianLayers.Conv1dGroupNJ(28, 1, 3, 1)
        #self.bn_conv1 = nn.BatchNorm2d(32)
        self.conv2 = BayesianLayers.Conv2dGroupNJ(28, 64, 3, 1)
        #self.bn_conv2 = nn.BatchNorm2d(64)
        
        '''self.conv3 = BayesianLayers.Conv3dGroupNJ(64,400,kernel_size = 5)
        
        self.flatten = FlattenLayer(1024)'''
        self.fc1 = BayesianLayers.LinearGroupNJ(9216, 128, clip_var=0.04, cuda=FLAGS.cuda)
        self.fc2 = BayesianLayers.LinearGroupNJ(128, 84, cuda=FLAGS.cuda)
        self.fc3 = BayesianLayers.LinearGroupNJ(84, 10, cuda=FLAGS.cuda)
        # layers including kl_divergence
        self.kl_list = [self.conv1,self.conv2,self.fc1, self.fc2, self.fc3]

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        '''x = self.conv3(x)
        x = F.sigmoid(x)'''
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        output = F.log_softmax(x, dim=1)
        return output

I’m not sure, how BayesianLayers are defined, but it seems you are mixing 1d, 2d, and 3d versions of conv layers in your model.
Based on your input shape, I guess you should stick to the 2d version.