Pt Train RNN Model

Pt Train RNN Model trains a recurrent model - RNN, GRU, or LSTM - for sequence classification, and it has the most knobs of any trainer in the pack because recurrent nets have real quirks the node has to handle for you. The pack's text workflows use it with GloVe word embeddings or with SentencePiece tokenization, and on IMDB sentiment it lands around 85% validation accuracy for a plain RNN and ~87% for GRU/LSTM - solid numbers for the architecture.

It's part of ComfyUI-Pt-Wrapper, the ~200-node pack that brings PyTorch training into ComfyUI's graph. It's the sequence sibling of the transformer trainers: same overall shape - model, dataloader, optimizer, loss in; trained model and loss histories out - but with three extra inputs that exist to keep you from shooting yourself in the foot on recurrent data.

How it works

The loop is the usual one: forward, loss, backprop, optimizer step, per-epoch scheduler step. What's interesting is the padding handling, controlled by use_valid_token_mean:

• True (default): the node averages the RNN's outputs over non-zero tokens before the final layer. Because batched sentences are padded to equal length, most of the last tokens are padding noise - averaging only the real tokens keeps the classifier from being dominated by it.
• False: it uses only the last token's output. The docs warn this is "weak if there's heavy padding." If your loss isn't decreasing, this flag is the pack's own recommended first check.

The other knob is linear_head. If your model already has a linear head on top of the recurrent layer (Ptn RNN Linear), the node calls it directly and computes loss on its output. Without a head, the node reshapes the raw recurrent output (rnn_output_reshape) into per-token predictions aligned with the labels. And if you want to control where the sequence starts, h_0 lets you pass an initial hidden state - shaped [num_layers, batch_size, hidden_size], or [2 * num_layers, batch_size, hidden_size] for bidirectional models. A PtZeros tensor is the natural neutral starting point here.

Inputs and outputs

• model, train_loader, optimizer, loss_function - the core. The text workflows pair this with Ptn BCE with Logits Loss.
• linear_head (BOOLEAN, true) and use_valid_token_mean (BOOLEAN, true) - the two sequence-specific switches above.
• epochs (INT, default 1), use_gpu, early_stopping (false), early_stopping_rounds (10), output_best_val_model (true), classification_metrics (false - flip it on to print validation accuracy).
• Optional: scheduler (PTLRSCHEDULER), val_loader, h_0 (TENSOR, initial hidden state).

Outputs: Model (PTMODEL), train loss, val loss tensors.

Installing the pack

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

Restart ComfyUI, or ComfyUI Manager → search "ComfyUI-Pt-Wrapper" → install (on the Comfy Registry; dependencies include pandas, scikit-learn, transformers, gensim and sentencepiece). Start from examples/workflows/rnn_classification.json (RNN), gru_classification.json, lstm_classification.json, or embedding_rnn_classification.json (training the embedding from scratch with the bundled SentencePiece tokenizer, vocab 32000).

Common issues

• Loss not decreasing. Check use_valid_token_mean first - heavy padding with the flag off is the pack's own documented failure mode.
• Head vs no-head mismatch. If linear_head is true but your model has no head (or vice versa), shapes won't line up at loss time. Match the flag to the model node you wired in.
• h_0 shape errors. Get the layer/batch/hidden dims right - bidirectional doubles the layer count. When in doubt, leave it unconnected and let the node use zeros.

It's the fiddliest trainer in the pack, but every knob exists because recurrent nets genuinely need it. Get use_valid_token_mean right and the rest is plain sailing.