updated the "Usage" section of README.md

README.md

@@ -40,11 +40,11 @@ tfms = v2.Compose([
     v2.ToDtype(torch.float32, scale=True)
 ])
 
+device = torch.device("cpu")
+
 import torch.nn as nn
 from huggingface_hub import PyTorchModelHubMixin
 
-device = torch.device("cpu")
-
 class MLP(nn.Module, PyTorchModelHubMixin):
     def __init__(self):
         super().__init__()
@@ -62,6 +62,12 @@ model = MLP.from_pretrained("sadhaklal/mlp-fashion-mnist")
 model.to(device)
 
 example = fashion_mnist['test'][0]
+
+import matplotlib.pyplot as plt
+
+plt.imshow(example['image'], cmap='gray')
+print(f"Ground truth: {id2label[example['label']]}")
+
 img = tfms(example['image'])
 x_batch = img.unsqueeze(0)
 
@@ -72,7 +78,7 @@ with torch.no_grad():
 proba = torch.softmax(logits, dim=-1)
 
 confidence, pred = proba.max(dim=-1)
-print(f"Predicted class: {pred[0].item()}")
+print(f"Predicted class: {id2label[pred[0].item()]}")
 print(f"Predicted confidence: {round(confidence[0].item(), 4)}")
 ```