tensorViz logo tensorViz / shape lab Runnable example

Your batch is fine.
Check the features.

A small image model reaches Linear with the wrong width. Change the input, follow the shapes, and find the contract that broke.

Shape arithmetic only. This page does not run Python, a model or inference.

Try an input

NCHW · 3 input channels

Batch 1–64. Height and width 2–256.

Choose the classifier
  1. 01Input
    [2, 3, 16, 16]
  2. 02Conv2d + ReLU3 → 8 channels · kernel 3 · padding 1
    [2, 8, 16, 16]
  3. 03MaxPool2d(2)Spatial sizes round down
    [2, 8, 8, 8]
  4. 04Flatten(start_dim=1)Keep the first dimension: batch
    [2, 512]
  5. 05Linear(128, 5)Expects 128 features per image
    512 ≠ 128

512 features arrived. The head expects 128.

8 channels × 8 height × 8 width = 512 features per image.

[2, 512]×[128, 5]incompatible

Linear stores its weight as [5, 128] and multiplies by its transpose.

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

Change the head to 512 to match this 16 × 16 example. Then try a wider image: the fixed width is a contract, not a universal fix.

A reshape can hide the error.

On the original feature map, reshape(-1, 128) asks PyTorch to infer a new first dimension. That dimension may no longer count your images.

This comparison uses the feature map before any optional adaptive pooling.

[2, 512]→ reshape →[8, 128]

2 images became 8 rows. Linear can now return [8, 5], but that is not one prediction per image.

Make the contract explicit.

Use Flatten(start_dim=1) to preserve the batch. Match the classifier width to the feature extractor and intended input size. Adaptive pooling is another design choice: it changes the computation and needs a different classifier.

Matching shapes only makes these operations compatible. It does not establish prediction quality or numerical equivalence.

Read and run the PyTorch example Explore recorded models in tensorViz