Batch size mismatch error when using loss function during the training of CNN

Hi,
I am training a simple CNN model. During the training I am getting an error message “ValueError: Expected input batch_size (16) to match target batch_size (32).”

I am using Python version 3.10.12.

This is my model:

# Create a convolutional neural network
class FashMNISTModelV2(nn.Module):
    # Model replicated called TinyVGG
    def __init__(self, input_shape: int, hiddent_units: int, output_shape:int):
        super().__init__()
        self.conv_block_1 = nn.Sequential(
            # Create a conv layer
            nn.Conv2d(in_channels=input_shape,
                      out_channels=hiddent_units,
                      kernel_size=3,
                      stride=1,
                      padding=1), # values we can set ourselfs are called hyperparamters
            nn.ReLU(),
            nn.Conv2d(in_channels=hiddent_units,
                      out_channels=hiddent_units,
                      kernel_size=3,
                      stride=1,
                      padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv_block_2 = nn.Sequential(
            # Create a conv layer
            nn.Conv2d(in_channels=hiddent_units,
                      out_channels=hiddent_units,
                      kernel_size=3,
                      stride=1,
                      padding=1), # values we can set ourselfs are called hyperparamters
            nn.ReLU(),
            nn.Conv2d(in_channels=hiddent_units,
                      out_channels=hiddent_units,
                      kernel_size=3,
                      stride=1,
                      padding=1),
            nn.ReLU(), 
            nn.MaxPool2d(kernel_size=2)
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=hiddent_units * 7 *7,
                      out_features=output_shape)
                     )
        
    def forward(self, x):
        x = self.conv_block_1(x)
        print("Output shape of conv block1: ", x.shape)
        x = self.conv_block_2(x)
        print("Output shape of conv block2: ", x.shape)
        x = self.classifier(x)
        print("Output shape of classifier: ", x.shape)
        return X

This is the training:

def train_step(model      : torch.nn.Module,
               dataloader : torch.utils.data.DataLoader,
               loss_fn    : torch.nn.CrossEntropyLoss,
               optimizer  : torch.optim.SGD,
               evaulation_function):
    epochs = 3
    for epoch in tqdm(range(epochs)):

        print("Epoch# : ", epoch)
        train_loss, train_acc = 0,0
        for batch, (X, y) in enumerate(dataloader):
            X = X.to("cuda")
            y = y.to("cuda")

            model.train()
            _y_pred = model(X)# forward pass

            loss = loss_fn(_y_pred, y) # calculate loss per batch
            train_loss += loss # accumatively add up the loss per epoch
            train_acc  += evaulation_function(y,# y_true values
                                              _y_pred.argmax(dim=1) # logits -> prediction labels 
                                              )

            optimizer.zero_grad() # optimizer zero grad

            loss.backward() # back propogation step

            optimizer.step()# optimizer step 

            # Print out how many samples have been seen
            if batch % 400 == 0:
                print(f"Looked at {batch * len(X)}/{len(dataloader.dataset)} samples")
        
        train_loss /= len(dataloader)
        train_acc  /= len(dataloader)

          ## Print out what's happening
        print(f"\nTrain loss: {train_loss:.5f} | Training accuracy: {train_acc}\n")

VS code is showing the error occurring at the loss function during the training function. My loss and optimizer functions are

loss_function = nn.CrossEntropyLoss()
optimizer     = torch.optim.SGD(params=model_2.parameters(),
                                lr=0.1)

Your model looks correct and keeps the batch size equal:

model = FashMNISTModelV2(1, 10, 10)
x = torch.randn(16, 1, 28, 28)
out = model(x)
# Output shape of conv block1:  torch.Size([16, 10, 14, 14])
# Output shape of conv block2:  torch.Size([16, 10, 7, 7])
# Output shape of classifier:  torch.Size([16, 10])

print(out.shape)
# torch.Size([16, 10])

However, note that you have a typo in your forward method in return X (capital X) while the activation is defined as x.

Thanks man. Not sure how I missed that typo.