PyTorch Tensor Shapes

Tensor shape is the number of elements in each dimension of a tensor. Shape errors are the most common cause of bugs in PyTorch code. A clear understanding of how shapes work — and how to change them — removes most of that frustration.

Reading a Tensor Shape

Think of shape like the dimensions of a physical object. A piece of paper has 2 dimensions: height and width. A stack of papers adds a third: depth. Tensors follow the same logic.

import torch

t0 = torch.tensor(42.0)                    # scalar
t1 = torch.tensor([1.0, 2.0, 3.0])        # vector
t2 = torch.tensor([[1.0, 2.0], [3.0, 4.0]])  # matrix
t3 = torch.rand(2, 3, 4)                  # 3D tensor

print(t0.shape)   # torch.Size([])
print(t1.shape)   # torch.Size([3])
print(t2.shape)   # torch.Size([2, 2])
print(t3.shape)   # torch.Size([2, 3, 4])

Shape Diagram

t2 = [[1, 2],    ← shape [2, 2]
      [3, 4]]
       │  └── 2 values per row
       └────── 2 rows

t3 = random tensor, shape [2, 3, 4]
  Visualize as 2 sheets of paper,
  each 3 rows × 4 columns

  Sheet 0:             Sheet 1:
  [· · · ·]            [· · · ·]
  [· · · ·]            [· · · ·]
  [· · · ·]            [· · · ·]

The Batch Dimension

When training a neural network, you don't feed one sample at a time — you feed a batch of samples. The first dimension of a tensor typically represents the batch size.

Single image shape:   (3, 224, 224)   ← channels, height, width
Batch of 32 images:   (32, 3, 224, 224)   ← batch + channels + height + width

# Adding a batch dimension to a single image tensor
single_image = torch.rand(3, 224, 224)    # shape: [3, 224, 224]
batched = single_image.unsqueeze(0)        # shape: [1, 3, 224, 224]
print(batched.shape)   # torch.Size([1, 3, 224, 224])

Reshaping Tensors

Reshaping changes the structure of a tensor without changing the data inside it. The total number of elements must stay the same.

t = torch.arange(12)            # 12 numbers: [0, 1, ..., 11]
print(t.shape)                  # torch.Size([12])

t_2d = t.reshape(3, 4)          # 3 rows × 4 columns = 12 elements
print(t_2d.shape)               # torch.Size([3, 4])

t_3d = t.reshape(2, 2, 3)       # 2 × 2 × 3 = 12 elements
print(t_3d.shape)               # torch.Size([2, 2, 3])

Using -1 in reshape

You can use -1 as a wildcard in reshape. PyTorch calculates the correct size for that dimension automatically.

t = torch.arange(24)
t2 = t.reshape(4, -1)    # 4 rows, PyTorch figures out columns: 24/4 = 6
print(t2.shape)           # torch.Size([4, 6])

t3 = t.reshape(-1, 8)    # columns = 8, rows = 24/8 = 3
print(t3.shape)           # torch.Size([3, 8])

Flattening a Tensor

Flattening converts any multi-dimensional tensor into a 1D vector. You use this before a fully connected layer in a CNN.

img = torch.rand(3, 28, 28)         # color image: channels × H × W
flat = img.flatten()                 # 3 × 28 × 28 = 2352 elements
print(flat.shape)                    # torch.Size([2352])

Checking Compatibility Before Operations

Two tensors can be added or multiplied only when their shapes are compatible. This diagram shows common shape errors and fixes:

a: shape [3, 4]
b: shape [3, 4]   → a + b  ✓ works (same shape)

a: shape [3, 4]
b: shape [4]      → a + b  ✓ works (broadcasting adds row dim to b)

a: shape [3, 4]
b: shape [3, 5]   → a + b  ✗ fails (column count mismatch)

Broadcasting Rules

Broadcasting lets PyTorch operate on tensors of different shapes by automatically expanding the smaller tensor. The rule: dimensions are compared from the right. Each pair must be either equal, or one of them must be 1.

a = torch.ones(3, 1)    # shape: [3, 1]
b = torch.ones(1, 4)    # shape: [1, 4]
c = a + b               # result shape: [3, 4]
print(c.shape)           # torch.Size([3, 4])

Broadcasting Visualization:
a:  [[1]      →  expand →  [[1, 1, 1, 1],
     [1]                    [1, 1, 1, 1],
     [1]]                   [1, 1, 1, 1]]

b:  [[1, 1, 1, 1]]  →  expand →  [[1, 1, 1, 1],
                                   [1, 1, 1, 1],
                                   [1, 1, 1, 1]]

a + b = [[2, 2, 2, 2],
          [2, 2, 2, 2],
          [2, 2, 2, 2]]  ← shape: [3, 4]

Common Shape Mistakes

Mismatched Matrix Multiplication

For matrix multiplication A @ B, the number of columns in A must equal the number of rows in B.

A = torch.rand(3, 4)    # 3 × 4
B = torch.rand(4, 5)    # 4 × 5
C = A @ B               # valid: 3 × 5

# This fails:
# B2 = torch.rand(3, 5)
# A @ B2   ← inner dims: 4 ≠ 3  → RuntimeError

Summary

Tensor shape describes the number of elements in each dimension. The first dimension is usually the batch size. Use reshape to change shape while keeping the same data — the total element count must stay equal. Use -1 in reshape to let PyTorch calculate one dimension automatically. Flatten collapses a multi-dimensional tensor to 1D. Broadcasting allows operations between tensors of different shapes by expanding smaller dimensions. Checking shapes before operations prevents the most common PyTorch errors.