PyTorch CNN Shape Errors: Conv2d Output Sizes, Channels, Flatten Bugs and How to Debug Them

AI, PyTorch
Page content

PyTorch: Zero to Hero — Step 06

CNN code is usually easy to write.

CNN shape bugs are usually easy to create.

A typical failure looks like this:

RuntimeError: Given groups=1, weight of size [32, 3, 3, 3],
expected input[64, 224, 224, 3] to have 3 channels,
but got 224 channels instead

Or this:

RuntimeError: mat1 and mat2 shapes cannot be multiplied

Or worse: the model runs, but the dimensions are wrong in a way that silently damages the architecture.

This post is a programmer-focused guide to the tensor geometry behind convolutional networks.

We will cover:

  • Conv2d input and output shapes
  • NCHW vs NHWC
  • channels-first vs channels-last memory format
  • calculating convolution output dimensions
  • stride, padding and dilation
  • pooling output sizes
  • flatten mistakes
  • Linear feature-count mismatches
  • AdaptiveAvgPool2d
  • debugging every intermediate activation
  • forward hooks
  • grouped and depthwise convolutions
  • model input contracts
  • building a complete CNN without guessing any dimensions

The goal is simple:

If a CNN shape breaks, you should be able to identify the exact layer and explain why.

1. The tensor contract for Conv2d

For normal batched image input, PyTorch Conv2d expects:

[N, C, H, W]

where:

N = batch size
C = channels
H = height
W = width

Example:

import torch

x = torch.randn(64, 3, 224, 224)

print(x.shape)

torch.Size([64, 3, 224, 224])

That represents:

64 RGB images
3 channels
224 pixels high
224 pixels wide

Now define a convolution:

import torch.nn as nn

conv = nn.Conv2d(
    in_channels=3,
    out_channels=32,
    kernel_size=3,
    stride=1,
    padding=1,
)

out = conv(x)
print(out.shape)

torch.Size([64, 32, 224, 224])

Notice what changed:

input:  [64,  3, 224, 224]
output: [64, 32, 224, 224]

The convolution changed the channel count from 3 to 32.

Because we used:

kernel_size=3
stride=1
padding=1

height and width stayed the same.

We can visualise this transformation as a pipeline step:

    flowchart LR
    A["Input [N, 3, H, W]"] --> B["Conv2d(3→32, k=3, s=1, p=1)"]
    B --> C["Output [N, 32, H, W]"]

2. The most common CNN mistake: NHWC data passed to NCHW code

A lot of image libraries represent images as:

[H, W, C]

and batches as:

[N, H, W, C]

PyTorch convolution layers expect:

[N, C, H, W]

Suppose your batch is:

x = torch.randn(64, 224, 224, 3)

and your layer is:

conv = nn.Conv2d(3, 32, kernel_size=3, padding=1)

This fails because PyTorch interprets dimension 1 as the channel dimension:

[N, C, H, W]
[64, 224, 224, 3]
     ^^^

It thinks the image has 224 channels.

The fix is usually:

x = x.permute(0, 3, 1, 2)

Now:

print(x.shape)

[64, 3, 224, 224]

A defensive helper is worth having:

def ensure_nchw(x: torch.Tensor) -> torch.Tensor:
    if x.ndim != 4:
        raise ValueError(f"Expected 4D image batch, got {tuple(x.shape)}")

    # Typical NHWC RGB batch
    if x.shape[-1] in (1, 3, 4) and x.shape[1] not in (1, 3, 4):
        x = x.permute(0, 3, 1, 2)

    return x

Use it at an input boundary, not randomly throughout the model.

Shape normalization belongs at the edge of the system.

3. Dimension order and memory format are not the same thing

This is subtle and worth understanding.

These two concepts are different:

dimension order
memory layout

PyTorch convolution APIs still conceptually operate on tensors shaped as:

[N, C, H, W]

But tensors can use a channels-last memory format internally.

Example:

x = torch.randn(16, 3, 224, 224)

print(x.shape)
print(x.stride())

Then:

x_cl = x.to(memory_format=torch.channels_last)

print(x_cl.shape)
print(x_cl.stride())

The shape remains:

[16, 3, 224, 224]

but the strides change.

Do not confuse:

x.permute(0, 2, 3, 1)

with:

x.to(memory_format=torch.channels_last)

The first changes dimension order.

The second changes physical memory layout while preserving logical dimension order.

Useful inspection:

def describe_layout(x: torch.Tensor) -> None:
    print("shape:", tuple(x.shape))
    print("stride:", x.stride())
    print("contiguous:", x.is_contiguous())
    print(
        "channels_last:",
        x.is_contiguous(memory_format=torch.channels_last)
        if x.ndim == 4
        else "n/a",
    )

4. Calculate Conv2d output size instead of guessing

For one spatial dimension, convolution output size is:

out = floor((in + 2p - d(k - 1) - 1) / s + 1)

where:

in = input size
k  = kernel size
s  = stride
p  = padding
d  = dilation

Write it once:

import math


def conv_out_size(
    size: int,
    kernel_size: int,
    stride: int = 1,
    padding: int = 0,
    dilation: int = 1,
) -> int:
    return math.floor(
        (
            size
            + 2 * padding
            - dilation * (kernel_size - 1)
            - 1
        ) / stride
        + 1
    )

Examples:

print(conv_out_size(224, 3, stride=1, padding=1))
print(conv_out_size(224, 3, stride=2, padding=1))
print(conv_out_size(32, 5, stride=1, padding=0))

Expected:

224
112
28

Now verify against PyTorch:

x = torch.randn(1, 3, 224, 224)

conv = nn.Conv2d(
    3,
    32,
    kernel_size=3,
    stride=2,
    padding=1,
)

y = conv(x)

print(y.shape)

[1, 32, 112, 112]

5. Tuple parameters matter

kernel_size, stride, padding and dilation can each be tuples.

Example:

conv = nn.Conv2d(
    3,
    16,
    kernel_size=(5, 3),
    stride=(2, 1),
    padding=(2, 1),
)

x = torch.randn(8, 3, 100, 200)
y = conv(x)

print(y.shape)

Now height and width follow different transformations.

A reusable 2D helper:

def pair(v):
    return (v, v) if isinstance(v, int) else v


def conv2d_output_hw(
    h: int,
    w: int,
    kernel_size,
    stride=1,
    padding=0,
    dilation=1,
):
    kh, kw = pair(kernel_size)
    sh, sw = pair(stride)
    ph, pw = pair(padding)
    dh, dw = pair(dilation)

    oh = conv_out_size(h, kh, sh, ph, dh)
    ow = conv_out_size(w, kw, sw, pw, dw)

    return oh, ow

Test it:

print(
    conv2d_output_hw(
        100,
        200,
        kernel_size=(5, 3),
        stride=(2, 1),
        padding=(2, 1),
    )
)

6. Why padding=1 often appears with kernel_size=3

For:

kernel=3
stride=1
dilation=1

using:

padding=1

preserves the spatial dimensions.

Example:

for size in [16, 32, 64, 224]:
    x = torch.randn(1, 3, size, size)
    conv = nn.Conv2d(3, 8, 3, padding=1)
    y = conv(x)
    print(size, "->", y.shape[-2:])

This is why blocks like this are common:

nn.Conv2d(64, 64, kernel_size=3, padding=1)

The layer changes representation without shrinking the feature map.

7. Stride is a spatial downsampler

Compare:

x = torch.randn(1, 3, 224, 224)

conv1 = nn.Conv2d(3, 32, 3, stride=1, padding=1)
conv2 = nn.Conv2d(3, 32, 3, stride=2, padding=1)

print(conv1(x).shape)
print(conv2(x).shape)

Output:

[1, 32, 224, 224]
[1, 32, 112, 112]

A stride-2 convolution often replaces explicit pooling in modern architectures.

8. Dilation changes the effective receptive field

A dilated kernel spaces kernel elements apart.

The effective kernel size is:

effective = dilation * (kernel_size - 1) + 1

Helper:

def effective_kernel(kernel_size: int, dilation: int) -> int:
    return dilation * (kernel_size - 1) + 1


for d in [1, 2, 3, 4]:
    print(d, effective_kernel(3, d))

1 3
2 5
3 7
4 9

This matters when you calculate output dimensions manually.

9. Pooling creates the same class of shape bugs

Example:

pool = nn.MaxPool2d(kernel_size=2, stride=2)

x = torch.randn(8, 32, 64, 64)
y = pool(x)

print(y.shape)

[8, 32, 32, 32]

A common CNN pattern is:

    flowchart LR
    A[Input] --> B[Conv2d]
    B --> C[ReLU]
    C --> D[MaxPool2d]
    D --> E[Conv2d]
    E --> F[ReLU]
    F --> G[MaxPool2d]
    G --> H[Flatten]
    H --> I[Linear]
    I --> J[Output]

The flatten step is where mistakes surface.

The real bug is usually several layers earlier.

10. The classic Linear shape mismatch

Consider:

class BrokenCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(16 * 32 * 32, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = torch.relu(x)
        x = self.pool(x)
        x = torch.flatten(x, 1)
        return self.fc(x)

This works for:

64 x 64 input

because pooling gives:

32 x 32

But feed it:

model = BrokenCNN()
x = torch.randn(8, 3, 224, 224)
model(x)

and the flattened feature count is now:

16 * 112 * 112

not:

16 * 32 * 32

The final layer is hardcoded to one image size.

11. Never debug a flatten mismatch from the error message alone

Instrument the forward pass:

class DebugCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(16 * 32 * 32, 10)

    def forward(self, x):
        print("input   ", tuple(x.shape))

        x = self.conv1(x)
        print("conv1   ", tuple(x.shape))

        x = torch.relu(x)
        print("relu    ", tuple(x.shape))

        x = self.pool(x)
        print("pool    ", tuple(x.shape))

        x = torch.flatten(x, 1)
        print("flatten ", tuple(x.shape))

        x = self.fc(x)
        print("fc      ", tuple(x.shape))

        return x

Run:

model = DebugCNN()
model(torch.randn(8, 3, 224, 224))

You immediately see the actual dimensions.

12. Use forward hooks when you do not want to modify the model

For existing models, hooks are cleaner.

def shape_hook(name):
    def hook(module, inputs, output):
        in_shapes = [
            tuple(x.shape)
            for x in inputs
            if isinstance(x, torch.Tensor)
        ]

        if isinstance(output, torch.Tensor):
            out_shape = tuple(output.shape)
        else:
            out_shape = type(output).__name__

        print(
            f"{name:<30} "
            f"{module.__class__.__name__:<20} "
            f"in={in_shapes} out={out_shape}"
        )

    return hook

Attach it:

handles = []

for name, module in model.named_modules():
    if name:
        handles.append(
            module.register_forward_hook(shape_hook(name))
        )

Run one batch:

with torch.no_grad():
    model(torch.randn(2, 3, 64, 64))

Then remove hooks:

for handle in handles:
    handle.remove()

This is one of the most useful debugging tools for unfamiliar PyTorch models.

13. Build a reusable shape tracer

from contextlib import contextmanager


@contextmanager
def trace_shapes(model: nn.Module):
    handles = []

    def make_hook(name):
        def hook(module, inputs, output):
            input_shapes = [
                tuple(t.shape)
                for t in inputs
                if isinstance(t, torch.Tensor)
            ]

            if isinstance(output, torch.Tensor):
                output_desc = tuple(output.shape)
            elif isinstance(output, (list, tuple)):
                output_desc = [
                    tuple(t.shape) if isinstance(t, torch.Tensor) else type(t).__name__
                    for t in output
                ]
            else:
                output_desc = type(output).__name__

            print(
                f"{name or '<root>':<35} "
                f"{module.__class__.__name__:<22} "
                f"{input_shapes} -> {output_desc}"
            )

        return hook

    for name, module in model.named_modules():
        if len(list(module.children())) == 0:
            handles.append(
                module.register_forward_hook(make_hook(name))
            )

    try:
        yield
    finally:
        for handle in handles:
            handle.remove()

Usage:

with trace_shapes(model):
    model(torch.randn(2, 3, 64, 64))

Now you have something reusable in real projects.

14. Remove hardcoded spatial dimensions with adaptive pooling

Instead of guessing what height and width reach your classifier, collapse them intentionally.

class RobustCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()

        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )

        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Linear(64, num_classes)

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

Try multiple input sizes:

model = RobustCNN()

for size in [32, 64, 128, 224, 256]:
    x = torch.randn(4, 3, size, size)
    y = model(x)
    print(size, "->", tuple(y.shape))

Every input produces:

[4, 10]

The classifier no longer depends on a fixed spatial resolution. We can visualise this with a quick scatter plot:

import matplotlib.pyplot as plt

sizes = [32, 64, 128, 224, 256]
shapes = []
for s in sizes:
    x = torch.randn(1, 3, s, s)
    with torch.no_grad():
        shapes.append(model.features(x).shape)

plt.plot(sizes, [sh[-1] for sh in shapes], marker='o')
plt.xlabel('Input size')
plt.ylabel('Feature map size (H/W)')
plt.title('AdaptiveAvgPool2d makes output independent of input resolution')
plt.grid(True)
plt.show()

15. Why AdaptiveAvgPool2d((1, 1)) is so useful

Suppose the feature tensor is:

[N, 256, 14, 14]

Adaptive global average pooling converts it to:

[N, 256, 1, 1]

Then:

x = torch.flatten(x, 1)

produces:

[N, 256]

Your classifier becomes:

nn.Linear(256, num_classes)

instead of:

nn.Linear(256 * 14 * 14, num_classes)

That reduces parameters and removes a large source of shape fragility.

16. Another option: infer the flatten dimension programmatically

Sometimes you want to preserve spatial features.

You can infer the flattened size with a dummy forward pass.

class InferredCNN(nn.Module):
    def __init__(self, input_shape=(3, 64, 64), num_classes=10):
        super().__init__()

        self.features = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )

        with torch.no_grad():
            dummy = torch.zeros(1, *input_shape)
            out = self.features(dummy)
            flat_features = out.flatten(1).shape[1]

        self.classifier = nn.Linear(flat_features, num_classes)

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

This is safer than handwritten arithmetic when the architecture changes frequently.

But the model still assumes the declared input resolution.

Adaptive pooling is more flexible when variable input size is desired.

17. flatten() has a batch-dimension trap

This is wrong:

x = x.flatten()

because it destroys the batch dimension too.

For a tensor:

[32, 64, 8, 8]

this becomes one long vector.

Usually you want:

x = x.flatten(1)

which produces:

[32, 4096]

Alternative:

x = torch.flatten(x, start_dim=1)

Or:

x = x.reshape(x.shape[0], -1)

I prefer:

x.flatten(1)

for CNN classifiers because the intent is obvious.

18. view() after permute() can fail

Consider:

x = torch.randn(4, 3, 32, 32)

x = x.permute(0, 2, 3, 1)

The tensor is now non-contiguous in the default memory format.

This may fail:

x.view(4, -1)

Safer options:

x = x.reshape(4, -1)

or:

x = x.contiguous().view(4, -1)

Inspect before guessing:

print(x.shape)
print(x.stride())
print(x.is_contiguous())

19. Grayscale input is not RGB input

A model defined as:

nn.Conv2d(3, 32, 3)

expects 3 channels.

MNIST-style image data may arrive as:

[N, 1, H, W]

Then you need:

nn.Conv2d(1, 32, 3)

or explicit conversion to three channels if the pretrained architecture requires RGB.

For example:

x_rgb = x.repeat(1, 3, 1, 1)

But do not do that blindly.

If you own the first layer, changing its in_channels may be cleaner.

20. Alpha channels can break RGB models too

Some image inputs have four channels:

RGBA

A model expecting RGB will fail.

Inspect inputs:

print(batch.shape)

If you deliberately want to discard alpha:

batch = batch[:, :3]

Again: normalize input semantics before the model, not deep inside it.

21. Grouped convolutions have additional contracts

Example:

conv = nn.Conv2d(
    in_channels=16,
    out_channels=32,
    kernel_size=3,
    padding=1,
    groups=4,
)

For grouped convolution, channel counts must be compatible with groups.

A useful assertion:

def validate_groups(in_channels, out_channels, groups):
    if in_channels % groups != 0:
        raise ValueError(
            f"in_channels={in_channels} must be divisible by groups={groups}"
        )

    if out_channels % groups != 0:
        raise ValueError(
            f"out_channels={out_channels} must be divisible by groups={groups}"
        )

22. Depthwise convolution is grouped convolution

A depthwise convolution uses:

groups = in_channels

Example:

channels = 32

conv = nn.Conv2d(
    channels,
    channels,
    kernel_size=3,
    padding=1,
    groups=channels,
)

Every channel is convolved independently.

A common depthwise-separable block is:

class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()

        self.depthwise = nn.Conv2d(
            in_channels,
            in_channels,
            kernel_size=3,
            padding=1,
            groups=in_channels,
            bias=False,
        )

        self.pointwise = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=1,
            bias=False,
        )

    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

Test:

block = DepthwiseSeparableConv(32, 64)
x = torch.randn(8, 32, 56, 56)
y = block(x)

print(y.shape)

[8, 64, 56, 56]

23. Build a CNN with explicit shape contracts

Here is a compact classifier that validates its input assumptions.

class ImageClassifier(nn.Module):
    def __init__(self, in_channels=3, num_classes=10):
        super().__init__()

        self.in_channels = in_channels

        self.features = nn.Sequential(
            nn.Conv2d(in_channels, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),

            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
        )

        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Linear(128, num_classes)

    def _validate_input(self, x):
        if x.ndim != 4:
            raise ValueError(
                f"Expected [N,C,H,W], got shape={tuple(x.shape)}"
            )

        if x.shape[1] != self.in_channels:
            raise ValueError(
                f"Expected {self.in_channels} channels at dim=1, "
                f"got shape={tuple(x.shape)}"
            )

        if x.shape[-2] < 4 or x.shape[-1] < 4:
            raise ValueError(
                f"Spatial dimensions too small: {tuple(x.shape[-2:])}"
            )

    def forward(self, x):
        self._validate_input(x)

        x = self.features(x)
        x = self.global_pool(x)
        x = x.flatten(1)
        return self.classifier(x)

Now failures happen near the boundary with a useful message.

That is much better than allowing a cryptic convolution error several calls later.

24. Test multiple image resolutions

model = ImageClassifier()
model.eval()

for h, w in [
    (32, 32),
    (64, 64),
    (128, 96),
    (224, 224),
    (256, 320),
]:
    x = torch.randn(2, 3, h, w)

    with torch.no_grad():
        y = model(x)

    print((h, w), "->", tuple(y.shape))

Because of adaptive pooling, each case returns:

[2, 10]

25. Build an activation auditor

Shapes are not the only thing worth checking.

You can inspect activations for NaNs, infinities and suspicious scale.

def activation_audit_hook(name):
    def hook(module, inputs, output):
        if not isinstance(output, torch.Tensor):
            return

        x = output.detach()

        finite = torch.isfinite(x)

        print(
            f"{name:<30} "
            f"shape={tuple(x.shape)!s:<20} "
            f"finite={finite.all().item()} "
            f"mean={x.float().mean().item(): .4e} "
            f"std={x.float().std().item(): .4e} "
            f"min={x.float().min().item(): .4e} "
            f"max={x.float().max().item(): .4e}"
        )

    return hook

Attach to convolution layers:

handles = []

for name, module in model.named_modules():
    if isinstance(module, nn.Conv2d):
        handles.append(
            module.register_forward_hook(
                activation_audit_hook(name)
            )
        )

Then:

with torch.no_grad():
    model(torch.randn(2, 3, 224, 224))

This starts bridging shape debugging into numerical debugging.

26. Parameter-count sanity checks

A CNN that suddenly has 150 million parameters when you expected 2 million probably contains a flatten/classifier mistake.

Use:

def count_parameters(model):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(
        p.numel()
        for p in model.parameters()
        if p.requires_grad
    )
    return total, trainable


total, trainable = count_parameters(model)

print("total:", f"{total:,}")
print("trainable:", f"{trainable:,}")

Now compare two heads.

Huge flattened head:

nn.Linear(256 * 14 * 14, 1000)

Global-pooled head:

nn.Linear(256, 1000)

The difference is enormous.

27. Calculate classifier parameter cost explicitly

def linear_parameter_count(in_features, out_features, bias=True):
    total = in_features * out_features
    if bias:
        total += out_features
    return total

Compare:

flat = linear_parameter_count(256 * 14 * 14, 1000)
gap = linear_parameter_count(256, 1000)

print(f"flatten head: {flat:,}")
print(f"global pool:  {gap:,}")

If your CNN parameter count explodes at the classifier, inspect the spatial dimensions before the linear layer.

28. Create a shape-contract decorator for development

For application code, explicit checks can save time.

def assert_image_batch(
    x: torch.Tensor,
    *,
    channels: int | None = None,
):
    if not isinstance(x, torch.Tensor):
        raise TypeError(f"Expected Tensor, got {type(x).__name__}")

    if x.ndim != 4:
        raise ValueError(
            f"Expected 4D [N,C,H,W], got {tuple(x.shape)}"
        )

    if channels is not None and x.shape[1] != channels:
        raise ValueError(
            f"Expected C={channels}, got shape={tuple(x.shape)}"
        )

    if not torch.is_floating_point(x):
        raise TypeError(
            f"Expected floating image tensor, got dtype={x.dtype}"
        )

Use:

assert_image_batch(x, channels=3)

This is especially useful after custom datasets and augmentation pipelines.

29. Unit-test tensor geometry

Model shape contracts deserve tests.

import pytest


@pytest.mark.parametrize(
    "shape",
    [
        (1, 3, 32, 32),
        (4, 3, 64, 64),
        (2, 3, 128, 96),
        (8, 3, 224, 224),
    ],
)
def test_classifier_shapes(shape):
    model = ImageClassifier(num_classes=7)
    model.eval()

    x = torch.randn(*shape)

    with torch.no_grad():
        y = model(x)

    assert y.shape == (shape[0], 7)

Test bad channel counts too:

def test_classifier_rejects_wrong_channels():
    model = ImageClassifier(in_channels=3)
    x = torch.randn(4, 1, 64, 64)

    with pytest.raises(ValueError):
        model(x)

CNN shape correctness should not depend on someone remembering to manually inspect one batch.

30. A complete tiny training example

Now put the geometry inside a real training loop.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset


torch.manual_seed(0)

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

num_classes = 5

x_train = torch.randn(1024, 3, 32, 32)
y_train = torch.randint(0, num_classes, (1024,))

train_loader = DataLoader(
    TensorDataset(x_train, y_train),
    batch_size=64,
    shuffle=True,
)

model = ImageClassifier(
    in_channels=3,
    num_classes=num_classes,
).to(device)

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
)

criterion = nn.CrossEntropyLoss()

for epoch in range(3):
    model.train()

    total_loss = 0.0
    total_examples = 0

    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad(set_to_none=True)

        logits = model(images)
        loss = criterion(logits, labels)

        loss.backward()
        optimizer.step()

        batch_size = images.shape[0]
        total_loss += loss.item() * batch_size
        total_examples += batch_size

    print(
        f"epoch={epoch} "
        f"loss={total_loss / total_examples:.4f}"
    )

The dataset is random, so do not expect meaningful accuracy.

The point is that all of the tensor contracts now connect:

    flowchart TD
    A["Dataset: [N,3,32,32]"] --> B["Conv2d"]
    B --> C["Feature maps"]
    C --> D["AdaptiveAvgPool2d"]
    D --> E["[N,128]"]
    E --> F["Linear"]
    F --> G["[N, num_classes]"]
    G --> H["CrossEntropyLoss"]

31. Debugging checklist: Conv2d expected X channels but got Y

When you see:

expected input to have 3 channels, but got 224 channels

check:

print(x.shape)

Then inspect:

Is the input NCHW?
Is it NHWC?
Is grayscale being passed to an RGB model?
Does the image contain alpha?
Did a previous layer output the expected channel count?
Does groups=... divide the channels correctly?

Do not change in_channels just to make the exception disappear.

First determine what the tensor actually represents.

32. Debugging checklist: mat1 and mat2 shapes cannot be multiplied

When the error appears immediately after flattening, print:

print("before flatten:", x.shape)
x = x.flatten(1)
print("after flatten:", x.shape)
print("fc expects:", self.fc.in_features)

Then ask:

Did image resolution change?
Did stride change?
Did padding change?
Was another pooling layer added?
Was a convolution removed?
Did channel count change?
Did flatten destroy the batch dimension?

The Linear layer is often where the mismatch is detected.

It is not necessarily where the mismatch was introduced.

33. Debugging checklist: model works for one resolution but not another

You probably hardcoded spatial dimensions into the classifier.

Search for code like:

nn.Linear(128 * 7 * 7, 1000)

Possible fixes:

1. enforce a fixed input resolution
2. calculate the feature shape explicitly
3. infer the flattened feature count
4. use adaptive pooling

Choose based on the model’s intended contract.

Do not blindly choose adaptive pooling if spatial position itself is part of the output semantics.

34. Debugging checklist: view size is not compatible...

Inspect:

print(x.shape)
print(x.stride())
print(x.is_contiguous())

If the tensor became non-contiguous through permute, use:

x.reshape(...)

or:

x.contiguous().view(...)

Do not scatter .contiguous() everywhere without understanding why it became necessary.

35. Debugging checklist: GPU memory unexpectedly explodes

Inspect shapes first.

A shape mistake can multiply activation memory dramatically.

For each major activation:

def tensor_megabytes(x: torch.Tensor) -> float:
    return x.numel() * x.element_size() / 1024**2

Hook version:

def memory_hook(name):
    def hook(module, inputs, output):
        if isinstance(output, torch.Tensor):
            print(
                f"{name:<30} "
                f"shape={tuple(output.shape)!s:<20} "
                f"size={tensor_megabytes(output):.2f} MiB"
            )
    return hook

A feature map of:

[64, 512, 224, 224]

is vastly more expensive than:

[64, 512, 14, 14]

Spatial downsampling is not cosmetic.

36. A compact CNN diagnostic utility

Put the common checks together:

def diagnose_cnn(
    model: nn.Module,
    input_shape=(2, 3, 224, 224),
    device="cpu",
):
    device = torch.device(device)
    model = model.to(device)
    model.eval()

    x = torch.randn(*input_shape, device=device)

    print("input")
    print("  shape:", tuple(x.shape))
    print("  dtype:", x.dtype)
    print("  device:", x.device)
    print()

    handles = []

    def make_hook(name):
        def hook(module, inputs, output):
            if not isinstance(output, torch.Tensor):
                return

            print(
                f"{name:<35} "
                f"{module.__class__.__name__:<20} "
                f"shape={tuple(output.shape)!s:<20} "
                f"dtype={str(output.dtype):<15} "
                f"device={output.device}"
            )

        return hook

    for name, module in model.named_modules():
        if name and len(list(module.children())) == 0:
            handles.append(
                module.register_forward_hook(make_hook(name))
            )

    try:
        with torch.no_grad():
            output = model(x)

        print()
        print("output:", tuple(output.shape))

    finally:
        for handle in handles:
            handle.remove()

Usage:

diagnose_cnn(
    ImageClassifier(num_classes=10),
    input_shape=(2, 3, 224, 224),
)

This is the sort of utility worth keeping in a real ML repository.

37. What a CNN actually changed from the previous posts

So far our networks mostly looked like:

features → Linear → activation → Linear

Images introduce a spatial structure:

channel
height
width

A fully connected layer immediately destroys that locality.

Convolution preserves it while sharing the same kernel weights across spatial positions.

That gives us tensors like:

[N, C, H, W]

through most of the feature extractor.

The central programming skill is therefore not merely knowing nn.Conv2d.

It is being able to track how those four dimensions evolve through the network.

38. The shape ledger technique

For non-trivial architectures, write down the expected shape transitions.

    flowchart TD
    A["Input [N, 3, 224, 224]"] --> B["Conv 3→32, stride 2 [N, 32, 112, 112]"]
    B --> C["Conv 32→64, stride 2 [N, 64, 56, 56]"]
    C --> D["Conv 64→128, stride 2 [N, 128, 28, 28]"]
    D --> E["Conv 128→256, stride 2 [N, 256, 14, 14]"]
    E --> F["AdaptiveAvgPool [N, 256, 1, 1]"]
    F --> G["Flatten [N, 256]"]
    G --> H["Linear [N, classes]"]

Then verify it with hooks.

This turns architecture debugging into contract verification.

39. Prefer assertions at semantic boundaries

Inside every convolution block, asserting every tensor dimension creates noise.

Better boundaries include:

dataset → model
backbone → classifier
encoder → decoder
feature extractor → detection head

Example:

features = self.backbone(x)

if features.ndim != 4:
    raise RuntimeError(
        f"Backbone contract broken: {tuple(features.shape)}"
    )

if features.shape[1] != 256:
    raise RuntimeError(
        f"Expected 256 backbone channels, got {features.shape[1]}"
    )

This is ordinary defensive programming applied to tensor code.

40. The shortest practical CNN debugging workflow

When a vision model breaks, do this before rewriting layers:

    flowchart TD
    A[Print raw batch shape] --> B[Verify NCHW semantics]
    B --> C[Verify dtype and device]
    C --> D[Run one batch only]
    D --> E[Attach forward shape hooks]
    E --> F[Locate first unexpected shape]
    F --> G[Inspect kernel/stride/padding/dilation at that layer]
    G --> H[Verify channel counts and groups]
    H --> I[Inspect flatten/classifier boundary]
    I --> J[Test multiple expected input resolutions]
    J --> K{Memory explodes?}
    K -- Yes --> L[Print activation sizes in MiB]
    K -- No --> M{Model runs but behaves strangely?}
    M -- Yes --> N[Audit activation mean/std/min/max; check for NaN/Inf]
    M -- No --> O[Done]

This is much faster than reading a 150-line stack trace and guessing.

Challenge: deliberately break the network

Take this model:

model = ImageClassifier(
    in_channels=3,
    num_classes=10,
)

Break it five different ways.

Break 1: NHWC input

x = torch.randn(8, 64, 64, 3)

Identify the failure and repair it at the input boundary.

Break 2: grayscale input

x = torch.randn(8, 1, 64, 64)

Decide whether the model or the data contract should change.

Break 3: remove adaptive pooling

Replace it with a hardcoded flatten head and test multiple resolutions.

Break 4: add a stride-2 convolution

Predict every downstream spatial dimension before running the model.

Break 5: permute then view

Create a non-contiguous tensor and inspect its strides before fixing it.

The objective is not merely to make the exception disappear.

The objective is to explain the tensor geometry that caused it.

Where the series goes next

We now have:

Step 00 — What Are We Actually Doing?
Step 01 — Tensor Shapes and Broadcasting
Step 02 — Autograd Debugging
Step 03 — Neural Network From Scratch
Step 04 — nn.Module Registration and state_dict
Step 05 — DataLoader Performance
Step 06 — CNN Shape Debugging

The next major jump is attention.

And again, we are not going to start with a giant Transformer class.

We are going to start with the part that programmers constantly get wrong:

batch
sequence
heads
embedding dimensions
attention masks
Q/K/V shapes

So Step 07 will be:

PyTorch Attention Shapes: Q, K, V, Multi-Head Attention Masks and Transformer Dimension Errors.

Once tensor shape reasoning is solid, attention becomes much less mysterious.