I am trying to implement masking in the “attention-weight” so that I update only the weight of the image label. Like there are 10 classes in CIFAR10. If the image is of class “5”, then I would like to update the weight of only class “5” by making its corresponding weight to 1 and the rest of the weights as “0”.