def show_predict_digit(model, dataset, index):
    image_tensor = dataset.data[index].unsqueeze(0).unsqueeze(0).to(torch.float32)
    plt.figure(figsize=(4, 4))
    plt.imshow(image_tensor.squeeze(0).squeeze(0))

    with torch.no_grad():
        output = model(image_tensor)

    return torch.argmax(output)