The output layer should have the number of classes as out_features.
Currently your output layer only returns one neuron, which corresponds to class0.
For a binary use case, this should work:
batch_size = 5
nb_classes = 2
in_features = 10
model = nn.Linear(in_features, nb_classes)
criterion = nn.CrossEntropyLoss()
x = torch.randn(batch_size, in_features)
target = torch.empty(batch_size, dtype=torch.long).random_(nb_classes)
output = model(x)
loss = criterion(output, target)
loss.backward()
However, this doesn’t seem to be the error you are seeing here.
As you can see in my example, target should be of type torch.long. Try to fix the shapes and call target = target.long() to transform the data type.
Alternatively, you could return just one neuron and use nn.BCEWithLogitsLoss as your criterion.
This would also work with a float target.