Multi-Class Cross Entropy Loss function implementation in PyTorch

I guess this might be the problem.
I would suggest to try the suggestion in this post to use 10 binary target images.