Pt Predict Classification Model

PtPredictClassificationModel is the "deploy" end of this pack's training story. Once you've trained a classifier, this node runs it on a fresh input and hands back three things: the predicted class name, the class id, and the probability the model assigned to that prediction. It's the node that turns "I trained a model" into "let me actually use it," and it's the sibling of the pack's Pt Predict Regression Model for the classification case.

Why you'd reach for it

Training a model in ComfyUI-Pt-Wrapper is only half the loop. The example workflows (dog-vs-cat, RNN text classification) follow the pattern: train with Pt Train Classification Model, save with Pt Save Model, then load the saved model back and run inference with this node. It does the real-world bookkeeping you don't want to hand-roll in the graph - adds a batch dimension to your single input, runs the model in inference mode, applies softmax, picks the argmax, and maps the winning class id back to a human-readable name.

Inputs

• model (PTMODEL) - your trained classifier.
• inputs (TENSOR) - one example to classify, without a batch dimension. The node adds it for you.
• class_id_to_name_map (PYDICT) - a Python dictionary mapping class ids to names, like {0: "cat", 1: "dog"}. If this is missing a key your model outputs, the node will crash on the lookup - it assumes every id it can produce has a name.
• use_gpu - off by default; flip on to move the model and input to CUDA for faster inference.

Outputs

Three sockets come out: class_name (STRING), class_id (INT), and probability (FLOAT). The probability is the softmax confidence for the winning class, between 0 and 1. These wire into the pack's text/display nodes, a Pt Show Text, or conditionals in your graph.

What's worth knowing

The model runs in inference mode and no gradients are computed, so this node is cheap and safe to call repeatedly. Note the pack deliberately uses train(False) rather than eval() in this path due to a ComfyUI security-checker quirk - behaviorally it's the same for the models you're likely to build. The one real footgun is the map: if your model outputs 2 classes but your class_id_to_name_map only names one, expect a KeyError exactly when you least want it. And use_gpu defaults to False, so if inference feels slow, that's the first toggle to flip.

Installing it

Part of ComfyUI-Pt-Wrapper (HowToSD's no-code PyTorch pack, a spin-off of ComfyUI-Data-Analysis). ComfyUI Manager → search "ComfyUI-Pt-Wrapper", or:

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

Restart after; first boot is slow while pandas, scikit-learn, transformers, sentencepiece, peft and friends install. No model files needed - but you'll want the pack's example workflows (and possibly the sibling ComfyUI-Data-Analysis pack for their display nodes) to run the full pipeline.