Introduction

Input shape is one of the most frequent sources of PyTorch errors because every layer expects dimensions in a particular order. If one axis is missing, swapped, or flattened at the wrong time, the model will either crash immediately or learn from the wrong representation.

The safest way to work in PyTorch is to document the expected shape at each boundary, inspect one batch from the data loader, and make reshaping operations explicit inside the model.

Learn the Common Shape Conventions

PyTorch usually puts the batch dimension first, but the remaining dimensions depend on the model family.

Typical conventions are:

• linear layers: batch, features
• image models with Conv2d: batch, channels, height, width
• sequence batches in many custom models: batch, sequence_length, features

A simple shape check makes this concrete:

1import torch
2
3vectors = torch.randn(32, 64)
4images = torch.randn(8, 3, 224, 224)
5sequence = torch.randn(16, 50, 128)
6
7print(vectors.shape)
8print(images.shape)
9print(sequence.shape)

The batch dimension is the first number in each case.

Match Layers to the Tensor They Expect

A frequent bug is connecting convolution output directly to a linear layer without flattening. Convolution layers output 4D tensors, while nn.Linear expects the last meaningful structure to be flattened into features.

1import torch
2import torch.nn as nn
3
4class ConvClassifier(nn.Module):
5    def __init__(self):
6        super().__init__()
7        self.features = nn.Sequential(
8            nn.Conv2d(3, 8, kernel_size=3, padding=1),
9            nn.ReLU(),
10            nn.MaxPool2d(2),
11        )
12        self.classifier = nn.Linear(8 * 16 * 16, 10)
13
14    def forward(self, x):
15        x = self.features(x)
16        x = torch.flatten(x, start_dim=1)
17        return self.classifier(x)
18
19
20model = ConvClassifier()
21out = model(torch.randn(4, 3, 32, 32))
22print(out.shape)

The flatten call keeps the batch dimension and collapses the rest into feature columns.

Avoid Hardcoding More Than Necessary

Hardcoded feature counts are brittle when the input resolution changes. One common improvement is to use adaptive pooling so the classifier sees a fixed-size tensor even when the image size varies.

1import torch
2import torch.nn as nn
3
4class SaferConvClassifier(nn.Module):
5    def __init__(self):
6        super().__init__()
7        self.features = nn.Sequential(
8            nn.Conv2d(3, 16, kernel_size=3, padding=1),
9            nn.ReLU(),
10            nn.AdaptiveAvgPool2d((1, 1)),
11        )
12        self.classifier = nn.Linear(16, 2)
13
14    def forward(self, x):
15        x = self.features(x)
16        x = torch.flatten(x, start_dim=1)
17        return self.classifier(x)

This reduces shape-related maintenance work because the classifier input size no longer depends on image width and height.

Validate One Batch Before Training

Many “model input shape” bugs actually start in the data pipeline. Before blaming the network, inspect one batch from the DataLoader.

for batch_x, batch_y in train_loader:
    print(batch_x.shape, batch_y.shape)
    break

This single check often reveals:

• missing channel dimensions
• wrong dtype or layout
• labels batched incorrectly
• unexpected padding or sequence length shape

If the batch is wrong here, the model code may be perfectly fine.

Add Assertions While Debugging

Shape assertions inside forward can save time during integration.

1class CheckedMLP(nn.Module):
2    def __init__(self):
3        super().__init__()
4        self.fc = nn.Linear(32, 16)
5
6    def forward(self, x):
7        assert x.dim() == 2, f"expected 2D input, got {x.shape}"
8        assert x.size(1) == 32, f"expected 32 features, got {x.size(1)}"
9        return self.fc(x)

These checks should not replace understanding, but they make failures immediate and readable.

Common Pitfalls

A common mistake is forgetting that Conv2d uses channels-first tensors. If your data arrives as height, width, channels, you need to reorder it.

Another mistake is flattening away the batch dimension by accident. torch.flatten(x) without start_dim=1 can collapse the whole tensor into one vector.

People also hardcode linear input sizes and then later change image size, pooling, or sequence length without updating the classifier.

Finally, do not assume the model is wrong before printing a real batch from the loader. Shape bugs often begin in transforms, collate functions, or preprocessing code.

Summary

• PyTorch models expect specific tensor shapes, and the batch dimension is usually first.
• Match each layer to the shape it expects, especially at conv-to-linear boundaries.
• Use torch.flatten(..., start_dim=1) carefully so the batch dimension stays intact.
• Print one real batch from the data loader before debugging the model internals.
• Shape assertions and adaptive pooling can make models much easier to maintain.