The shape of out is expected to be [batch_size, nb_classes], while yours seems to be only [batch_size]. If you are dealing with a binary classification use case, you could use nn.BCEWithLogitsLoss (or nn.BCELoss, if you already applied sigmoid on your output).