How can i visualize all the predictions with tensorboard in Pytorch?

Question

I tried to visualize the actual images with the predictions to figure out how my algorithm is performing and which labels are being predicted incorrectly. But when I set the step, during visualization on the tensorboard, it does not show all the steps. So, it does not show all the training images and their labels. Instead, I can only see a few examples out of all the training images.

writer = SummaryWriter(log_dir='graphs')

def matplotlib_imshow(img):
    npimg = img.cpu().numpy()
    npimg = np.transpose(npimg, (1, 2, 0))
    plt.imshow((npimg * 255).astype(np.uint8))

def images_to_probs(net, images):
    output = net(images)
    _, preds_tensor = torch.max(output, 1)
    preds = np.squeeze(preds_tensor.cpu().numpy())
    return preds, [F.softmax(el, dim=0)[i].item() for i, el in zip(preds, output)]

def plot_classes_preds(net, images, labels):
    preds, probs = images_to_probs(net, images)
    fig = plt.figure(figsize=(6, 6))
    for idx in np.arange(4):
        ax = fig.add_subplot(1, 4, idx+1, xticks=[], yticks=[])

        matplotlib_imshow(images[idx])
        ax.set_title("{0}, {1:.1f}%\n(label: {2})".format(
            classes[preds[idx]],
            probs[idx] * 100.0,
            classes[labels[idx]]),
                    color=("green" if preds[idx]==labels[idx].item() else "red"))
    return fig

The following is my training loop where I used the global step as step.

for epoch in range(epochs):
    epoch_start_time = time.time()
    losses = []
    total_batch_images = 0
    batch_correct_pred = 0
    step = 0
    #save model
    # if batch_accuracy>best_acc:
    #     best_acc = batch_accuracy
    #     checkpoint = {'state_dict': model.state_dict(),'acc' : batch_accuracy, 'epoch' : epoch,  'optimizer': optimizer.state_dict()}
    #     save_checkpoint(checkpoint)

    model.train()
    for batch_idx, (images, labels) in enumerate(train_loader):
        # Get data to cuda if possible
        
        images = images.to(device=device)
        labels = labels.to(device=device)

        # forward
        scores = model(images)
        loss = criterion(scores, labels)

        losses.append(loss.item())

        # backward
        optimizer.zero_grad()
        loss.backward()

        # gradient descent or adam step
        optimizer.step()

        # visualizing Dataset images
        # img_grid = torchvision.utils.make_grid(images)
        # writer.add_image('Xray_images', img_grid, global_step = step)

        # calculation running accuracy
        model.eval()
        _, predictions = scores.max(1)
        num_correct = (predictions == labels).sum()
        batch_correct_pred += float(num_correct)
        total_batch_images += predictions.size(0)

        writer.add_figure('predictions vs. actuals',
                           plot_classes_preds(model, images, labels),
                           global_step=step)      
        step += 1

I think that the problem is in the last line of writer.add_figure() where the global_step is defined. However, any help in this regard would be really appreciated.