Assertion `THIndexTensor_(size)(target, 0) == batch_size' failed

I know this has been asked but I could not apply the suggestions there to solve my problem. My code is a simple copy of Pytorch’s tutorial to classify CIFAR-10 dataset and looks like this:

from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
def init(self):
super(Net, self).init()
self.conv1 = nn.Conv2d(3, 32, 3)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(32, 64, 3)
self.conv3 = nn.Conv2d(64, 128, 3)
self.conv4 = nn.Conv2d(128, 256, 3)
self.conv5 = nn.Conv2d(256, 512, 3)
self.conv6 = nn.Conv2d(512, 1024, 3)
self.fc1 = nn.Linear(36864, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)

def forward(self, x):
    #x = self.pool(F.relu(self.conv1(x)))
    x = F.relu(self.conv1(x))
    #x = self.pool(F.relu(self.conv2(x)))
    x = F.relu(self.conv2(x))
    x = self.pool(F.relu(self.conv3(x)))
    x = F.relu(self.conv4(x))
    x = F.relu(self.conv5(x))
    x = self.pool(F.relu(self.conv6(x)))        
    x = x.view(-1, 36864)
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x
    #print (x.size())

net = Net()

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(2): # loop over the dataset multiple times

running_loss = 0.0
for i, data in enumerate(trainloader, 0):
    # get the inputs
    inputs, labels = data

    # wrap them in Variable
    inputs, labels = Variable(inputs), Variable(labels)

    # zero the parameter gradients
    optimizer.zero_grad()

    # forward + backward + optimize
    outputs = net(inputs)
    print(outputs.size())
    print(outputs)
    print(labels)
    print(labels.size())
    loss = criterion(outputs, labels)
    print(output.size())
    print(labels.size())
    loss.backward()
    optimizer.step()

    # print statistics
    running_loss += loss.data[0]
    if i % 2000 == 1999:    # print every 2000 mini-batches
        print('[%d, %5d] loss: %.3f' %
              (epoch + 1, i + 1, running_loss / 2000))
        running_loss = 0.0

print(‘Finished Training’)

and the error is this:

torch.Size([1, 10])
Variable containing:
1.00000e-02 *
-4.6679 -1.3744 0.1477 5.1275 -8.8241 -0.7043 -1.2692 -7.0015 7.4909 -9.7750
[torch.FloatTensor of size 1x10]

Variable containing:
4
1
1
9
[torch.LongTensor of size 4]

torch.Size([4])

RuntimeErrorTraceback (most recent call last)
in ()
18 print(labels)
19 print(labels.size())
—> 20 loss = criterion(outputs, labels)
21 print(output.size())
22 print(labels.size())

/home/jigyasa/anaconda2/lib/python2.7/site-packages/torch/nn/modules/module.pyc in call(self, *input, **kwargs)
222 for hook in self._forward_pre_hooks.values():
223 hook(self, input)
–> 224 result = self.forward(*input, **kwargs)
225 for hook in self._forward_hooks.values():
226 hook_result = hook(self, input, result)

/home/jigyasa/anaconda2/lib/python2.7/site-packages/torch/nn/modules/loss.pyc in forward(self, input, target)
480 _assert_no_grad(target)
481 return F.cross_entropy(input, target, self.weight, self.size_average,
–> 482 self.ignore_index)
483
484

/home/jigyasa/anaconda2/lib/python2.7/site-packages/torch/nn/functional.pyc in cross_entropy(input, target, weight, size_average, ignore_index)
744 True, the loss is averaged over non-ignored targets.
745 “”"
–> 746 return nll_loss(log_softmax(input), target, weight, size_average, ignore_index)
747
748

/home/jigyasa/anaconda2/lib/python2.7/site-packages/torch/nn/functional.pyc in nll_loss(input, target, weight, size_average, ignore_index)
670 dim = input.dim()
671 if dim == 2:
–> 672 return _functions.thnn.NLLLoss.apply(input, target, weight, size_average, ignore_index)
673 elif dim == 4:
674 return _functions.thnn.NLLLoss2d.apply(input, target, weight, size_average, ignore_index)

/home/jigyasa/anaconda2/lib/python2.7/site-packages/torch/nn/_functions/thnn/auto.pyc in forward(ctx, input, target, *args)
45 output = input.new(1)
46 getattr(ctx._backend, update_output.name)(ctx._backend.library_state, input, target,
—> 47 output, *ctx.additional_args)
48 return output
49

RuntimeError: Assertion `THIndexTensor_(size)(target, 0) == batch_size’ failed. at /opt/conda/conda-bld/pytorch_1503966894950/work/torch/lib/THNN/generic/ClassNLLCriterion.c:54

The batch size I’ve taken is 4. But I think the output of print(labels.size()) should be 10 ijstead of 4 because that is the total number of classes.
What am I missing?

Your input likely has batch size 1, which is not equal to the label with size 4.

If you have batch size 4, label should be LongTensor of size [4] representing the correct class indices.

I have loaded the data like this:

transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root=’./data’, train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root=’./data’, train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
shuffle=False, num_workers=2)

classes = (‘plane’, ‘car’, ‘bird’, ‘cat’,
‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’)

Here, the batch_size argument is 4. Isn’t everything in order then?

Hmm then problem might be elsewhere.

Are you sure that each image gives you 36864 values? Is it possible that this is 4x than it really needs to be?