Hello Ikram!



 ulhaqi12:

        x = F.softmax(self.layer5(x))
        
        return x
...
criterian = nn.CrossEntropyLoss()



For starters, get rid of the final softmax() activation and pass the
result of your last linear layer directly to CrossEntropyLoss.
As mentioned (but not really emphasized) in the CrossEntropyLoss 1
documentation, CrossEntropyLoss expects raw-score logits
(rather than the probabilities that are produced by softmax()).
Best.
K. Frank