I am learning to re-train Resnet, with one more layer of 9 classes on top of final 1000.
I think I do not understand criterion function.
criterian = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr = 0.0001, momentum = 0.9)
for epoch in range(20):
    running_loss  = 0.0
    for X, Y in training:
        optimizer.zero_grad()
        output = net(X)

        print(output)
        print(output.shape)
        print(Y)
        print(Y.shape)
        print(Y.argmax())
        loss = criterian(output, Y.argmax())
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(epoch, ':', running_loss)

It seems legit for me (since I don’t know yet how to create batches I am providing image one-by-one).
I have output of tensor of shape 1x9 and target is a one element tensor, which specifies into which class it should be targeting.
However the output is:
tensor([[0.1639, 0.2125, 0.1067, 0.1212, 0.1434, 0.0613, 0.0512, 0.0225, 0.1172]],
       device='cuda:0', grad_fn=<SoftmaxBackward>)
torch.Size([1, 9])
tensor([0., 0., 1., 0., 0., 0., 0., 0., 0.], device='cuda:0')
torch.Size([9])
tensor(2, device='cuda:0')
Traceback (most recent call last):
  File "im9_train.py", line 173, in <module>
    loss = criterian(output, Y.argmax())
  File "/home/szandala/imagenet/venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/szandala/imagenet/venv/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 961, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/szandala/imagenet/venv/lib/python3.8/site-packages/torch/nn/functional.py", line 2468, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/home/szandala/imagenet/venv/lib/python3.8/site-packages/torch/nn/functional.py", line 2260, in nll_loss
    if input.size(0) != target.size(0):

What am I doing wrong?