Find the predicted class for each row

For a batch of class scores, you usually want one predicted class per row. Calling argmax() without dim instead returns one index into the entire flattened tensor.

import torch

logits = torch.tensor([[1.0, 9.0], [7.0, 3.0]])
print(logits.argmax().item())          # 1
print(logits.argmax(dim=1).tolist())   # [1, 0]

The flattened values are [1, 9, 7, 3], whose largest value is at index 1. The row-wise result separately identifies the largest score in each row. A global index can exceed the number of classes, or happen to look like a valid class index while still answering the wrong question.

Choose the axis to reduce

For shape (batch, classes), reducing dimension 1 removes the class axis and leaves shape (batch,). This assumes the default keepdim=False. With keepdim=True, the result instead has shape (batch, 1).

Use dim=-1 to name the last axis. For this two-dimensional input, it means the same thing as dim=1. Naming the axis makes the intended grouping clear; omitting it is appropriate when you really want one global result.

Ties

If a row contains multiple equal largest values, argmax returns the index of the first one. This specifies tie handling for the supplied scores; it does not guarantee that an entire model produces identical scores across runs, devices or precision settings.

Give each tensor axis a meaning

A tensor's shape tells you how many entries each axis has. It does not tell you what those entries represent. For example, (batch, features) and (features, batch) use different axis orders. Name the axes before choosing an operation.

A reduction combines entries along an axis, for example by summing or averaging. Averaging the features of a (B, D) tensor gives one number per example, with shape (B,). Keeping the reduced axis gives shape (B, 1) instead.

For [[2, 4, 6], [1, 3, 5]], the row means with the axis kept are [[4], [3]]. Subtracting them from the input gives [[-2, 0, 2], [-2, 0, 2]]: each row has been centered using its own mean.

Reshape and transpose do different jobs

A reshape regroups elements using the operation's traversal order. A transpose changes the axis order. For the usual row-major reshape, [[1, 2, 3], [4, 5, 6]] reshaped to (3, 2) becomes [[1, 2], [3, 4], [5, 6]]. Transposing the original instead gives [[1, 4], [2, 5], [3, 6]].

Both results have shape (3, 2), but their entries mean different things. Check values as well as shapes. Unequal axis sizes make many axis mistakes easier to spot than square examples do.

Concatenate along an axis or stack on a new one

Concatenation extends an existing axis. Stacking creates a new axis. For two arrays of shape (2, 3), concatenating on axis zero gives (4, 3). Stacking on a new leading axis gives (2, 2, 3).