TrainModel

This is the node the whole pack is named for. TrainModel takes a model, a dataset (or features/labels tensors), and runs an actual PyTorch training loop inside the ComfyUI graph - epochs, batching, loss, optimizer steps, and a running metric - then hands you back the trained model, the per-epoch metrics, and a "best" model checkpoint. If you've ever wanted to train a neural network without leaving ComfyUI and without writing a for epoch in range(...) loop, this is the entire appeal of EternalKernel PyTorch Nodes in one box.

Context first: the training advice you'll find in the wider community - rank tables, LR schedules, text-encoder freezing - is all about diffusion LoRAs, and none of it applies here. This is classical supervised ML on whatever nn.Sequential you built: tiny nets, real epochs, honest gradients. MNIST in, classifier out. It's a learning tool and a prototyping tool, not a diffusion fine-tuner.

How it works

The loop is textbook PyTorch: it builds a DataLoader with your batch_size and shuffle=True, then for each epoch it zeroes gradients, forward-passes each batch, computes the loss, backprops, and steps the optimizer. Per epoch it records one metric value, and it tracks the lowest-loss epoch as the "best model," saving and reloading the model each epoch via torch.save/torch.load to do it.

The data input is flexible: give it a dataset, or a features tensor plus a labels tensor (it wraps those in a TensorDataset itself). Provide neither and it raises a clear error.

The inputs that matter

Most of the dropdowns are generated straight from your torch install, which is neat and occasionally overwhelming:

• epochs (default 1) - bump it. One epoch of MNIST is technically training and practically nothing.
• batch_size (default 1) - a batch of 1 is slow and noisy; 32–128 is the sane range for this pack's toy models.
• loss_function (default CrossEntropyLoss) - must match your problem. CrossEntropyLoss for classification with a class-count Linear head, MSELoss for regression. A mismatch produces shape errors you'll chase for a while.
• optimizer (default SGD) - Adam with its default LR (0.001) is usually the least-fiddly choice here.
• metrics (default accuracy) - the thing written into the metrics output each epoch. Accuracy for classification; loss functions also work as metrics.
• device (default torch.device('cpu')) - the big gotcha: it defaults to CPU even if you have a GPU. Set it to cuda explicitly or you'll wonder why your 3090 is watching.
• create_samples (default FALSE) - a quirky extra that dumps side-by-side prediction-vs-label image PNGs into the current working directory when TRUE. Cute, usually off.

The three outputs: model (the trained net), metrics (a LIST of one value per epoch), and best model (the lowest-loss epoch's weights).

Installing it

Shared with the whole pack:

cd ComfyUI/custom_nodes
git clone https://github.com/TashaSkyUp/EternalKernelPytorchNodes.git
cd EternalKernelPytorchNodes
pip install -r requirements.txt

Restart ComfyUI and it's under ETK/pytorch - or ComfyUI Manager, searching "EternalKernel PyTorch Nodes". The requirements add scipy, scikit-learn, transformers and friends; your ComfyUI Python env needs them (pip install -r does it).

Troubleshooting

Four traps account for most of the pain. Device: check it's cuda if you expect GPU speed. Loss shape: CrossEntropyLoss wants raw logits of shape (batch, num_classes) and integer labels; feed it a regression output and you get a matrix-shape error. Metrics are last-batch only - the value per epoch is computed on the final batch, not averaged over the whole epoch, so don't read too much into small wobbles. And freezing gets undone: TrainModel force-sets requires_grad = True on all parameters, so a SetModelTrainable freeze before this node is ignored. Also note best model is tracked by lowest total epoch loss and is pickled with torch.save - fine for your own runs, not something to share with strangers. If metrics come out looking like noise, reduce the learning rate or lower batch_size, same as everywhere.