I am getting the same error. please help me. Following is my code.
for epoch in range(1, args.num_epochs+1):
epoch_loss = []
    for step, (images, labels) in enumerate(loader):
        if args.cuda:
            images = images.cuda()
            labels = labels.cuda()

        inputs = Variable(images)
        targets = Variable(labels)
        outputs = model(inputs)
        optimizer.zero_grad()
        loss = criterion(outputs, targets[:, 0])
        loss.backward()
        optimizer.step()

        epoch_loss.append(loss.data[0])
        if args.steps_plot > 0 and step % args.steps_plot == 0:
            image = inputs[0].cpu().data
            image[0] = image[0] * .229 + .485
            image[1] = image[1] * .224 + .456
            image[2] = image[2] * .225 + .406
            board.image(image,
                f'input (epoch: {epoch}, step: {step})')
            board.image(color_transform(outputs[0].cpu().max(0)[1].data),
                f'output (epoch: {epoch}, step: {step})')
            board.image(color_transform(targets[0].cpu().data),
                f'target (epoch: {epoch}, step: {step})')
        if args.steps_loss > 0 and step % args.steps_loss == 0:
            average = sum(epoch_loss) / len(epoch_loss)
            print(f'loss: {average} (epoch: {epoch}, step: {step})')
        if args.steps_save > 0 and step % args.steps_save == 0:
            filename = f'{args.model}-{epoch:03}-{step:04}.pth'
            torch.save(model.state_dict(), filename)
            print(f'save: {filename} (epoch: {epoch}, step: {step})')