RuntimeError: size mismatch the model return tensor with shape (num_class)

1

this is the model with 6 classes as output

image
image869×445 67.2 KB

and this is the training loop

with batch size =2
the shape of the batch ; torch.Size([2, 3, 224, 224])
but in forward the model return tensor with shape of (6) instead of shape(2,6) [batch_size,class]
the error message :
RuntimeError: size mismatch (got input: [6], target: [2])

2

What is the shape of inp? Based on the code I guess inp = torch.tensor([kl1, kl2, kl3]) might create a tensor in an unexpected shape.

3

inp is a tensor contains 3 numbers