Pt Train Model

Pt Train Model is the general-purpose trainer in ComfyUI-Pt-Wrapper, and it's the one to reach for when your task doesn't fit the pack's more opinionated trainers. Where the classification nodes hard-code their loss, this one lets you plug in any loss function you want - which makes it the node for regression-style tasks, custom objectives, or just exploring how different losses change training behavior.

The wiring is a full training loop, no code required: connect a model (PTMODEL from one of the pack's model-building nodes), a train_loader (PTDATALOADER from a data loader node), an optimizer (PTOPTIMIZER - Adam, AdamW, SGD, etc.), and a loss_function (PTLOSS). Then set the dials that matter:

• epochs - how many passes over the data. Default 1, which is basically a warmup; you'll usually bump this.
• use_gpu - Boolean, default off. Yes, really. If you have a GPU, flip this or you'll train on CPU and wonder why it's glacial.
• early_stopping + early_stopping_rounds - with a val_loader connected, this watches validation loss and stops if it doesn't improve for N rounds. Default off.
• output_best_val_model - default true: with a val_loader, returns the model that had the best validation loss rather than the last epoch's weights.

It also accepts an optional scheduler (PTLRSCHEDULER, e.g. cosine annealing) to decay the learning rate over training. The outputs are Model, train loss, and val loss - the loss outputs are tensors containing the per-epoch histories, which you can feed into the pack's plotting/printing nodes to watch training happen.

A critical behavior to know: the node's IS_CHANGED always returns NaN, which means ComfyUI considers it permanently "changed" - every time you hit queue, it retrains, even with identical inputs. That's correct for training, but it means you must not leave a training node dangling in a workflow that queues repeatedly. And when output_best_val_model kicks in, the best checkpoint is written to the pack's models/tmp/best_model.pt before being returned.

Install

ComfyUI Manager (search "ComfyUI-Pt-Wrapper"), or:

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

Restart ComfyUI. Heavy requirements (transformers, datasets, peft, accelerate, scikit-learn, pandas, seaborn, matplotlib, gensim, sentencepiece). No model files to download.

Common issues

Forgetting use_gpu is the #1 "why is this slow" complaint. Shape mismatches between your loss function and your labels are the #2 - if the loss node and label tensors disagree, you'll get a terse PyTorch error, and the fix is usually a cast (labels should be int64 for classification-style losses) or a reshape. And remember this node re-trains on every queue, so disconnect it once you're happy with the weights.