how to make prediction in bayesian convolutional neural network using pyro and pytorch

Question

I'm trying to classify the sign language hand gesture images using Bayesian CNN, I made the following code:

class BCNN(PyroModule):
    def __init__(self):
        super(BCNN, self).__init__()

        self.conv1 = PyroModule[nn.Conv2d](in_channels=1, out_channels=28, kernel_size=(3,3), stride=1, padding=1)
        self.conv1.weight = PyroSample(dist.Normal(0., 1.).expand(self.conv1.weight.shape).to_event(self.conv1.weight.dim()))
        self.conv1.bias = PyroSample(dist.Normal(0., 1.).expand(self.conv1.bias.shape).to_event(self.conv1.bias.dim()))

        # rest of the codes
        
        self.softmax = nn.LogSoftmax(dim=1)

    def forward(self, x, y=None):
        
        # rest of the codes

        logits = self.softmax(output)
        
        with pyro.plate("data", x.shape[0]):
            obs = pyro.sample("obs", dist.Categorical(logits=logits), obs=y)
        return logits

#----------------------------------

model = bcnn.BCNN()
guide = AutoDiagonalNormal(model)
optim = pyro.optim.Adam({"lr": 0.03})
svi = SVI(model, guide, optim, loss=Trace_ELBO())

#----------------------------------

num_iterations = 5
loss = 0

for j in range(num_iterations):
    loss = 0
    for batch_id, (images, labels) in enumerate(train_loader):
        # calculate the loss and take a gradient step
        loss += svi.step(images, labels)
    normalizer_train = len(train_loader.dataset)
    total_epoch_loss_train = loss / normalizer_train
    
    print("Epoch ", j+1, " Loss ", total_epoch_loss_train)

#----------------------------------

def summary(samples):
    site_stats = {}
    for k, v in samples.items():
        site_stats[k] = {
            "mean": torch.mean(v.to(dtype=float), 0),
            "std": torch.std(v.to(dtype=float), 0),
        }
    return site_stats

#------------------------------------------
# input: image | output (class): label
#------------------------------------------

predictive = pyro.infer.Predictive(model=model, guide=guide, num_samples=10, return_sites=("obs", "_RETURN"))
samples = predictive(image)
pred_summary = summary(samples)

#----------------------------------

a = pred_summary["_RETURN"]
b = pred_summary["obs"]

print(f"Label: {label.item()}")
print(torch.max(a['mean'].view(1,-1), 1)[1].item())

but the point is that I don't know how to use the trained model to compute the output. I mean giving the input image and get the class of the model using the trained model? I'm using the above codes but It gives me the wrong output always.