In your current script the model output contains 40 elements ([20, 2]) while the target contains only 20 ([20, 1, 1]). Assuming the target contains the class indices, you could remove the unnecessary dimensions via target = target.squeeze(2).squeeze(1), make sure it’s a LongTensor, and use nn.CrossEntropyLoss instead. On the other hand, you could also one-hot encode the target and use nn.BCEWithLogitsLoss instead.