ComfyUI Node

Pt Argmax

The node that turns logits into a prediction

By HowToSD·Created about a year ago·Updated about a year ago· 7
Pt Argmax
  • tens
  • TENSOR
dim
keepdimfalse

Your classifier's last layer spits out a tensor of scores - one per class - and the answer is "the class with the highest score." Pt Argmax is the node that extracts that answer: it returns the indices of the maximum values, which for a classifier means the predicted class. In the pack's dog-vs-cat and CIFAR-10 examples, this is literally the node where a model stops being a tensor factory and starts making predictions.

It's a reduction op in ComfyUI-Pt-Wrapper (HowToSD's 200-node PyTorch pack, the spin-off of ComfyUI-Data-Analysis), the argmax twin of Pt Argmin. If you're computing accuracy - "did the model guess right?" - the pipeline is: model logits → argmax → predicted class indices → compare with the label tensor. Argmax is the piece that turns softmax/logits into a concrete decision.

How it works. Three inputs:

  • tens - the TENSOR to reduce.
  • dim - which axis to argmax over. Here's the pack's quirk: dim is a string field (multiline, so you can even type more than one dimension). Type it as plain text, e.g. 1 to argmax across the class axis of a (batch, classes) tensor. Leave it empty ("") and it flattens the whole tensor - that gives you one global max index, which is rarely what you want. Read it like an integer: 0, 1, 2, or -1 for the last axis.
  • keepdim - a boolean. True keeps the reduced axis as a size-1 dimension (so (8, 10)(8, 1)); False drops it ((8, 10)(8,)). Default False.

Output is a TENSOR of indices (integers), same rank handling as above. The mechanism is torch.argmax(tens, dim=dim, keepdim=keepdim) with the string parsed to an int - the pack validates it.

The classic mistakes. The dim string is the #1 trap: people expect a dropdown, get a text field, and either leave it blank (silently flattening to a single global index) or type a float (1.0), which the parser rejects. Get in the habit of typing a bare integer and double-checking you hit the class axis. For (batch, classes) data that's 1 (or -1). Second: argmax returns positions, not values - the number you get is "which class," not "how confident." Confidence lives in the raw scores before argmax; if you need it, split the stream before this node. Third: ties are broken deterministically but arbitrarily - if two classes have identical scores, you get the lower index, which can make accuracy look flaky on toy data.

Install: ComfyUI Manager → "ComfyUI-Pt-Wrapper", or:

cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper

then restart. No model downloads; the pack's heavy requirements.txt (transformers, sklearn, sentencepiece, pinned gensim) is the only install cost.

Troubleshooting: output is a single number when you expected a batch - dim was empty or wrong; you flattened. Parser error on dim - type a bare integer, no decimals. Accuracy all zeros - the predicted indices don't line up with your label indexing; check that class ordering matches your dataset's label order. If keepdim behavior confuses you, turn it on (True) - keeping the axis makes downstream comparison nodes happier.

CategoryData Analysis

Inputs (3)

NameTypeDefaultDescription
tensTENSOR
dimSTRING
keepdimBOOLEANfalse

Outputs (1)

NameTypeDescription
TENSORTENSOR