Hello Neo!



 neoSilex:


And the output dimension of the model is [No x Co x Ho x Wo]
where,
No -> is the batch size (same as Ni)
Co -> is the number of classes that the dataset have!
Ho -> the height of the image (which is the same as Hi in almost all cases)
Wo -> the width of the image (which is the same as Wi in almost all cases)




Ho and Wo must be the same as Hi and Wi in all cases, not
just in “almost all” cases.


number of channels=1 or 3
the target dimension is [Ni x Ci x Hi x Wi]
where,
Ni -> the batch size
Ci -> the number of channels (which is 3 or 1)
Hi -> the height of the image
Wi -> the width of the image 




This might be what you have, but it simply won’t work.
CrossEntropyLoss requires that, for a model output of shape
[No, Co, Ho, Wo], the target have shape [No, Ho, Wo]
(and that the values of the target are integer class labels that
run from 0 to Co - 1).
If your target has this extra “channel” dimension (Ci), it won’t
work (and Hi and Wi must match Ho and Wo, as well).  (Just
to be clear, you can’t have the Ci dimension at all, even if
Ci = 1.)
Good luck.
K. Frank