Hi @albanD. I actually figured out the problem. The softmax is giving same values sometimes. So when i run:
mask[tensor_c < max_value] = 0
mask[tensor_c >= max_value] = 1

I get this:
tensor([[0., 0., 0., 0., 1., 1., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 1., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 1., 0., 1., 1., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 1., 1., 0., 1., 0., 0.],
        [0., 0., 0., 0., 1., 1., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 1., 1., 0., 1., 1., 1., 1., 0., 0.],
        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 1., 1., 0., 0., 0., 0., 0.]])

Therefore,
The number of elements in mask2[mask2 == 1] would not be equal to the number of elements in values_batch which is the batch size (13 in this example), since there are more than one “1” for each batch.
Is there a clean way to solve this? For example, mask[tensor_c>=max_value] = 1 only inserts a “1” one time, and if the same value is repeated, it will just ignore placing a 1. I assume this happens at the beginning of training only, then softmax would be more accurate and produce a single max value.