title: ValueError: Expected input batch_size (8) to match target batch_size (32)
body: <p>The task being performed is simple classification on the MNIST dataset. The problem I'm facing is that for the architecture shown below, I'm able to perform the training step which was giving a similar error earlier but that was due to some architectural issue, but now it's giving the error for the test step. I've tried to change the batch size and checked the size of the inputs at each step, but couldn't understand what's going wrong.</p>
  <p>Thank you for your help!</p>
  <p><strong>LeNet5 architecture:</strong></p>
  <pre><code>class LeNet5(nn.Module):
  
  def __init__(self, n_classes):
      super(LeNet5, self).__init__()
      
      self.feature_extractor = nn.Sequential(            
          nn.Conv2d(in_channels=1, out_channels=6, kernel_size=4, stride=1),
          nn.Tanh(),
          nn.AvgPool2d(kernel_size=2),
          nn.Conv2d(in_channels=6, out_channels=16, kernel_size=4, stride=1),
          nn.Tanh(),
          nn.AvgPool2d(kernel_size=2),
          nn.Conv2d(in_channels=16, out_channels=120, kernel_size=4, stride=1),
          nn.Tanh()
      )
  
      self.classifier = nn.Sequential(
          nn.Linear(in_features=120*2*2, out_features=84),
          nn.Tanh(),
          nn.Linear(in_features=84, out_features=n_classes),
      )
  
  
  def forward(self, x):
    # print(x.shape)
    x = self.feature_extractor(x)
    # print(x.shape)
    x = x.view(-1, 120*2*2)
    # print(x.shape)
    logits = self.classifier(x)
    # print(logits.shape)
    return logits
  </code></pre>
  <p>Training function:</p>
  <pre><code>def train(model, epoch, delta=1e-5):
    global best_train_acc, best_test_acc
    device = 'cuda'
    model.to(device)
    print('\nEpoch: %d' % epoch)
    model.train()
    train_loss = 0
    correct = 0
    total = 0
    total_acc = 0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        
        # for params in model.feature_extractor.parameters():
        #   params.requires_grad = False
  
        # model.feature_extractor[6].parameters().requires_grad = True
        optimizer.zero_grad()
        outputs = model(inputs)
        print(outputs.shape)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
  
        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        epsilon, best_alpha = optimizer.privacy_engine.get_privacy_spent(delta)
  
        print(f'Epoch {epoch} Step {batch_idx}/{len(trainloader)}', 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                      % (train_loss/(batch_idx+1), 100.*correct/total, correct, total),  f&quot;(ε = {epsilon:.2f}, δ = {delta}) for α = {best_alpha}&quot;)
        
        acc = 100.*correct/total
        total_acc += acc
        if acc&gt;best_train_acc:
            best_train_acc = acc
      
    return total_acc/len(trainloader)
  </code></pre>
  <p>Test function:</p>
  <pre><code>def test(model, epoch, model_name, lr):
    global best_train_acc, best_test_acc
    device = 'cuda'
    model.eval()
    test_loss = 0
    correct = 0
    total = 0
    total_acc = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(testloader):
            # print(inputs.shape, inputs[0].shape)
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
  
            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
  
            print(f'Epoch {epoch} Step {batch_idx}/{len(testloader)}', 'Test Loss: %.3f | Test Acc: %.3f%% (%d/%d)'
                          % (test_loss/(batch_idx+1), 100.*correct/total, correct, total))
  
    # Save checkpoint.
    acc = 100.*correct/total
    total_acc += acc
    if acc &gt; best_test_acc:
        print('Saving..')
        state = {
            'model': model.state_dict(),
            'acc': acc,
            'epoch': epoch,
        }
        if not os.path.isdir('/content/drive/My Drive/checkit'):
            os.mkdir('/content/drive/My Drive/checkit')
        torch.save(state, f'/content/drive/My Drive/checkit/ckpt_{model_name}_{lr}_MNIST.pth')
        best_test_acc = acc
        
    return total_acc
  </code></pre>
  <p>Error</p>
  <pre><code>    ValueError                                Traceback (most recent call last)
  &lt;ipython-input-25-5c2bb9b3b3d1&gt; in &lt;module&gt;()
       22 for epoch in range(start_epoch, start_epoch+50):
       23     LeNet5_train_acc.append([train(model, epoch), epoch+1])
  ---&gt; 24     LeNet5_test_acc.append([test(model, epoch, 'LeNet5', lr), epoch+1])
       25 
       26 best_acc.append([best_train_acc, best_test_acc])
  
  4 frames
  /usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
     2260     if input.size(0) != target.size(0):
     2261         raise ValueError('Expected input batch_size ({}) to match target batch_size ({}).'
  -&gt; 2262                          .format(input.size(0), target.size(0)))
     2263     if dim == 2:
     2264         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
  
  ValueError: Expected input batch_size (8) to match target batch_size (32).
  </code></pre>