natelang:

when I print out labels in nested for loop it is all 1.


If that’s the case, your model would only have to output the highest logit for class1 and would achieve a perfect classification.
Since you are using a custom dataset, I would recommend to check its implementation again and make sure the real targets are returned.
Feel free to post the Dataset implementation in case you get stuck.