Thank you for your reply @xian_kgx .
And I think I understand what you’re saying.
But the reason I took argmax for input data is that I had to reshape the data size to fit the expected target size for nn.Crossentropyloss.
https://pytorch.org/docs/stable/nn.html#crossentropyloss 3
So again, input data’s shape to VAE is
(Channel=Class, Height, Width) = (20, 100, 100)
Output data’s from VAE is
(Channel=Class, Height, Width) = (20, 100, 100)
And I want to minimize the loss between those two.
To match the expected target shape for nn.Crossentropyloss, I took argmax.
But the problem is, I don’t want to take argmax for input data since I want to keep each classes information. (Not only class which has the highest probability)
Thank you