Fix “mat1 and mat2 shapes cannot be multiplied” in PyTorch

01 / REPRODUCE IT

Start with a small model you can inspect.

You need Python and PyTorch. These files were tested together on CPU with Python 3.12.14 and PyTorch 2.7.1. Use an environment with PyTorch already installed, or follow the official installation instructions for your system.

Save the files below in one folder. The models use random initial weights and synthetic inputs; they do not download a dataset or checkpoint.

From that folder, run:

python verify.py

python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
python verify.py

The script catches the deliberate errors, checks the corrections, and finishes with All CPU checks passed. Fix the feature contract; do not reshape away the mismatch. Here is the broken model:

"""A deliberately wrong Linear width after flattening a small image model."""

from torch import nn


class BrokenClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.flatten = nn.Flatten(start_dim=1)
        self.classifier = nn.Linear(128, 5)  # Wrong: a 16 x 16 input produces 512 features.

    def forward(self, x):
        x = self.features(x)
        x = self.flatten(x)
        return self.classifier(x)

For a batch of two 16 × 16 RGB images, the failing call reports:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x512 and 128x5)

02 / FOLLOW THE DIMENSIONS

The inner numbers have to agree.

In (2x512 and 128x5), the first matrix contains 2 examples with 512 features each. The second is the Linear weight used in the multiplication: 128 input features and 5 outputs. The mismatch is 512 ≠ 128; the batch size of 2 is fine.

PyTorch stores Linear(128, 5).weight with shape [5, 128] and uses its transpose in x @ weight.T + bias. That is why the error says 128x5. See the Linear input and weight contract.

The convolution keeps the spatial size because its 3 × 3 kernel uses padding 1. The 2 × 2 pooling reduces 16 × 16 to 8 × 8. Flattening from dimension 1 keeps the batch dimension and combines the remaining dimensions: 8 * 8 * 8 = 512.

When debugging your own model, inspect x.shape just before the failing Linear call and compare x.shape[-1] with layer.in_features. Check the image size, pooling and channel count that produced it before choosing a replacement width.

03 / REPAIR THE FEATURE CONTRACT

Give the head the 512 features it receives.

Keep the feature extractor and flatten operation. Correct the head for this 16 × 16 input:

"""The fixed-resolution correction: 8 channels times 8 times 8 spatial values."""

from torch import nn


class FixedClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.flatten = nn.Flatten(start_dim=1)
        self.classifier = nn.Linear(8 * 8 * 8, 5)

    def forward(self, x):
        x = self.features(x)
        x = self.flatten(x)
        return self.classifier(x)

The corrected forward call produces [2, 5]: five outputs for each of the two images. The verification script also checks batches 1 and 5, finite outputs, and agreement between batched and one-image-at-a-time inference.

Do not copy 512 into every model. It follows from this model’s channels and spatial size. Replacing a trained Linear layer also changes its parameter shape; an old 128-feature head cannot be loaded directly into the new 512-feature head.

04 / KEEP THE BATCH INTACT

A reshape can hide the error and break the model.

The tempting workaround is x.reshape(-1, 128). Our two flattened images contain 2 * 512 = 1024 values, so that expression produces [8, 128]. The broken head now runs and returns [8, 5]: eight predictions for two images.

The verification script reproduces that counterexample. Changing the batch dimension to make a matrix multiplication fit loses the one-row-per-image contract. Preserve it with Flatten(start_dim=1) (or torch.flatten(x, start_dim=1)) and fix the feature width where it is wrong. The flatten documentation explains which dimensions are combined.

05 / CHECK THE INPUT SIZE TOO

A fixed width does not accept every image size.

With a 16 × 20 or 20 × 16 input, this feature extractor produces 640 values per image. The corrected 512-feature head then fails with (2x640 and 512x5). The test includes both non-square cases so a successful 16 × 16 run does not hide that limit.

If the intended model should accept different spatial sizes, one architectural option is AdaptiveAvgPool2d((1, 1)) before flattening. It averages each channel’s spatial values, leaving eight features per image, followed by Linear(8, 5). PyTorch documents the pooling output-size contract.

"""A different architecture: spatial averaging before the classification head."""

from torch import nn


class PooledClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.flatten = nn.Flatten(start_dim=1)
        self.classifier = nn.Linear(8, 5)

    def forward(self, x):
        x = self.features(x)
        x = self.pool(x)
        x = self.flatten(x)
        return self.classifier(x)

This version passes the tested square and non-square inputs. It changes the computation and the head’s weights. Use it when spatial averaging fits your model design; it is not a way to preserve an existing fixed-resolution checkpoint’s predictions.