Tornike:

i had that 2 because i wanted to return tensor of 2 probabilities of 2 classes and to return higher probability of these 2 as a prediction.
again, what you advised, it worked but it is always predicting that ‘1’ class


Like I wrote in the above answer, I just assumed you had 1 class because your target tensor is of size torch.Size([10]).
If you want to have two classes than this is were your problem lies and not with the model.
If you have 2 classes, the target tensor would have to be of size torch.Size([10, 2])
Just make sure you really want 2 classes. If you want a binary classification were the final output is either True or False, than a one dimensional output with a probability between 0.0 and 1.0 (which you get by putting the modles output into a sigmoid function) is what most people would use.
If you still want to have 2 classes, than your first model was totally fine as it was.
You just need your target to have 2 classes (be 2 dimensional) as well