NCHW to NHWC, for free

Image tensors come in two conventions and you will convert between them for the rest of your life:

NCHW   (batch, channels, height, width)   what PyTorch layers expect
NHWC   (batch, height, width, channels)   what images and matplotlib expect

The conversion moves no data. It is a relabelling of axes, and a tensor’s axes live entirely in its shape and stride.

The free way

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

Read it as: axis 0 of the result is axis 0 of the input, axis 1 of the result is axis 2 of the input, and so on. Sizes and strides are reordered together, the buffer is untouched.

The expensive ways

Every one of these produces the same numbers and allocates a new buffer:

torch.stack([x[:, c] for c in range(C)], dim=-1)   # C slices, then a copy
x.permute(0, 2, 3, 1).contiguous()                 # a copy, on request
torch.from_numpy(x.numpy().transpose(0, 2, 3, 1))  # a round trip

The second is the one that actually happens in real code. .contiguous() gets added because something downstream complained once, and then it stays forever, copying every batch of every epoch.

Sometimes you genuinely need it — a few kernels require contiguous input, and view requires it. The rule is to add it at the point that needs it, not as punctuation after every permute.

Your task

def to_nhwc(x: torch.Tensor) -> torch.Tensor

Convert an NCHW tensor to NHWC without copying the buffer.

The starter is the .contiguous() version: right numbers, right shape, and a full copy of the batch every time it runs.