Hey, Extremely sorry for the late reply, I see you are getting a wrong prediction but a very good accuracy because you have not applied normalization. Due to this when you try to apply
torch.max(yb, dim = 1), it takes the maximum value from an unnormalized set of prediction and gets the result according to that. So the class with the highest value is the predicted value.
As far as I know the solution to this problem is applying normalization in the train and test data.
The below image explains how I came to know about the error.
I printed out the prediction you were getting and checked what is the maximum value, (In the below image I have already applied normalization). When normalization was not applied the value was not present in a particular range rather it was having a high range and was biased. After I applied normalization I was getting correct results from your code.
Results after normalization.