GridSearchTraining

Instead of hand-tuning epochs, batch_size, and loss_function one at a time, GridSearchTraining runs every combination you specify and reports back. It's the pack's automation layer over TrainModel: you give it a model, a param_grid, and your data, and it loops through the Cartesian product of the grid - every combo of every value - training a fresh run each time. Three outputs: best_model, best_params, and all_metrics for every combo tried.

Grid search is the blunt instrument of hyperparameter tuning - exhaustive, simple, and guaranteed to try stuff you'd never think of. For the small toy models this pack builds, it's genuinely a nice way to see how sensitive training is to batch_size or loss_function. Just don't confuse it with Bayesian optimization; it's brute force with a spreadsheet.

How it works

param_grid is a JSON string mapping hyperparameter names to lists of values, with a sensible default of {"epochs": [1, 2], "batch_size": [32, 64], "loss_function": ["MSELoss", "CrossEntropyLoss"]}. It parses the JSON, builds itertools.product of all value combinations, and for each combo constructs TrainModel kwargs (any extra optional inputs - dataset, features/labels tensors - get passed through to every run) and trains. Results accumulate in all_metrics.

The bug you need to know about

The "best model" selection is backwards. The comparison logic adds the current run's metrics to all_metrics before checking whether it's better than the minimum - so it's comparing the run against itself, and the "is it better" check can never fire. Net effect: best_model and best_params will always be the first combination tried, not the best one. Treat best_model/best_params as "the first run" and read all_metrics yourself to pick the winner. This is grounded in the source; it's not a rumor.

Inputs and outputs

• model (TORCH_MODEL) - the Sequential to train.
• param_grid (STRING) - JSON object of hyperparameter → list of values. Keys must match TrainModel's parameters (epochs, batch_size, loss_function, optimizer, ...).
• Optional: dataset (TORCH_DATASET), or features tensor + labels tensor (TORCH_TENSOR) - pass one or the other, same rules as TrainModel.

Outputs: best_model (TORCH_MODEL), best_params (DICT), all_metrics (LIST).

Install

ComfyUI Manager, search "EternalKernel PyTorch Nodes", or:

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

Restart ComfyUI; node under ETK/pytorch. No model files. Requirements are the standard stack plus scipy, scikit-learn, transformers, einops.

Common issues

• Trusting best_model. It's the first combo. Wire all_metrics into a text/plot node, read the losses, and pick the real winner's parameters yourself.
• Malformed JSON. The field is plain text; a missing brace or quote fails the json.loads. Keep it valid (the default string is a good template).
• Combinatorial explosion. 3 params × 3 values each = 27 full training runs. On CPU that's slow; keep the grid small.
• Wrong keys = ignored keys. Keys that TrainModel doesn't know get swallowed. epochs, batch_size, loss_function, optimizer are the safe ones.

The pack is small and tutorial-free, so this is the kind of thing you'd only discover by reading the source - and here it is, documented. Pack-wide quirk: it patches ComfyUI's validator to ignore return_type_mismatch, so check your wires. And when the grid finishes, remember: the winner is in all_metrics, not in the box labeled "best."