Skip to course content
Free PyTorch course

Advanced Deep Learning with PyTorch

Unit 02.00: Tensor shape, dtype, device, and batch dimension

Most PyTorch errors are shape errors, and most shape errors come from losing track of one dimension: the batch.

Four properties describe every tensor

Shape is how many elements sit along each axis. dtype is the numeric type. device is where the memory lives. And by convention the first axis is the batch — one entry per example.

Image batches are ordered *batch, channels, height, width*:

import torch

t = torch.zeros(8, 3, 32, 32)     # batch, channels, height, width
print("shape :", tuple(t.shape))
print("dtype :", t.dtype)         # float32 by default
print("device:", t.device)        # cpu

# The first dimension is almost always the batch. Losing it is the most
# common shape bug in PyTorch.
one = t[0]
print("single item shape       :", tuple(one.shape))          # (3, 32, 32)
print("restored batch dimension:", tuple(one.unsqueeze(0).shape))  # (1, 3, 32, 32)

# dtype matters: integer division and float division are different operations.
ints = torch.tensor([1, 2, 3])
print("int tensor mean fails:", end=" ")
try:
    ints.mean()
except RuntimeError as e:
    print(type(e).__name__, "- cast to float first")
print("as float:", ints.float().mean().item())

Indexing with t[0] drops the batch dimension and returns a single image of shape (3, 32, 32). Most layers will reject that, because they expect a batch. unsqueeze(0) puts it back.

The dtype example matters too: mean() on an integer tensor raises an error rather than silently rounding.

The mistake this prevents

Passing a single example to a model without restoring the batch dimension. The error message talks about matrix sizes, which sends people hunting through their layer definitions when the real fix is one unsqueeze(0).

Takeaway

Read the shape before you read the error. Batch first, and check the dtype when the operation involves division or an average.