Which axis disappeared

dim is the axis that disappears. That single sentence resolves most confusion about reductions, and the corollary is what this problem is about: a reduction with no dim makes every axis disappear.

m = torch.tensor([[1., 9.],
                  [7., 3.]])

m.argmax()          # tensor(1)      the flat index of the global maximum
m.argmax(dim=1)     # tensor([1, 0]) the index within each row

The first is not “the argmax of the matrix” in any per-row sense. It flattens to [1, 9, 7, 3], finds position 1, and hands you that. If your matrix has four columns, a returned index of 6 is row 1 column 2, and nothing tells you that.

The reason this is worth its own problem

argmax() returning a plausible small integer is exactly the shape of a per-row answer when there are few rows. Classifying a batch of two:

preds = logits.argmax()      # tensor(3)

A batch of 2 over 4 classes gives a flat index in 0..7, which looks like a class index and is not. It will index something successfully and be wrong.

The rule

Always pass dim to a reduction unless you genuinely mean “collapse the whole tensor to one number”. sum(), mean(), max(), argmax(), all(), any() all behave this way.

And when the result feeds a broadcast, pass keepdim=True, for the reasons the broadcasting track covered.

Ties

argmax returns the first maximal position. That is documented, and it is the reason two different runs of the same code agree.

Your task

def predicted_classes(logits: torch.Tensor) -> torch.Tensor

logits is (batch, classes). Return a 1-D tensor of length batch: the index of the largest logit in each row.