Pt Argmax
The node that turns logits into a prediction
- tens
- TENSOR
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- theTENSORto reduce.dim- which axis to argmax over. Here's the pack's quirk:dimis a string field (multiline, so you can even type more than one dimension). Type it as plain text, e.g.1to 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-1for the last axis.keepdim- a boolean.Truekeeps the reduced axis as a size-1 dimension (so(8, 10)→(8, 1));Falsedrops it ((8, 10)→(8,)). DefaultFalse.
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.
Inputs (3)
| Name | Type | Default | Description |
|---|---|---|---|
| tens | TENSOR | — | |
| dim | STRING | — | |
| keepdim | BOOLEAN | false | — |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| TENSOR | TENSOR | — |