Changing original label of MNISTto new label

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
import numpy as np

training_data = datasets.MNIST(root = “data”,train = True,download = True,transform = ToTensor())
test_data = datasets.MNIST(root = “data”,train = False,download = True,transform = ToTensor())

labels_map = {0: “Label : 0”,

          1: "Label : 1",

          2: "Label : 2",

          3: "Label : 3",

          4: "Label : 4",

          5: "Label : 5",

          6: "Label : 6",

          7: "Label : 7",

          8: "Label : 8",

          9: "Label : 9",

          }

figure = plt.figure(figsize=(8,8))

cols, rows = 3, 3

for i in range(1,cols*rows + 1):

sample_idx = torch.randint(len(training_data),size=(1,)).item()

img,label = training_data[sample_idx]

A = np.zeros(10)

A[label] = 1

print(A)

figure.add_subplot(rows, cols, i)

plt.title(labels_map[label])

plt.axis(“Off”)

plt.imshow(img.squeeze(),cmap = “gray”)

plt.show()

for this how to change label
if you select number ‘3’ image should come of containing as ‘1’ all other images should be ‘0’
if an original label is same as the selected number than change into ‘1’
if an original label is not same as the selected number than change into ‘0’