Just by skimming through the code, it seems you are using a wrong activation function for the criterion.
nn.BCELoss expects a sigmoid function applied on the model outputs.
However, you mentioned that his is not a multi-label classification, so I assume it’s a multi-class classification?
If that’s the case, remove the softmax/sigmoid and use nn.CrossEntropyLoss.