AddSoftmaxLayerNode

When your classifier in the EternalKernel PyTorch Nodes pack spits out raw numbers like [2.1, -0.4, 0.8], what are you supposed to do with that? A softmax. AddSoftmaxLayerNode appends torch.nn.Softmax (or LogSoftmax) to the end of your nn.Sequential and turns those raw logits into a set of values that sum to 1 - i.e. an actual probability distribution over your classes. The biggest number becomes the predicted class, and the others tell you how confident the model isn't.

You'll usually reach for it at the end of a classification network, after the last linear layer. One PyTorch-specific gotcha: if you train with CrossEntropyLoss (the TrainModel node's default), do not bolt a plain Softmax on the end first - cross-entropy in PyTorch already applies a log-softmax internally, and feeding it a softmax'd tensor is a double-softmax that trains badly. That's why log_softmax=True exists here: it gives you LogSoftmax, which is what you want when you're computing your own NLLLoss.

How it works

It builds nn.Softmax(dim=dim) - or nn.LogSoftmax(dim=dim) if you flip log_softmax - and inserts it at the end of your Sequential, mutating your model in place and returning it. That's the whole mechanism; the interesting decision is dim.

Inputs

• model (TORCH_MODEL) - the Sequential being extended. Start from SequentialModelProvider.
• dim - which axis to normalize over. For a batch of class scores shaped (batch, num_classes), dim=1 is what you want (normalize across classes, not across the batch). dim=0 normalizes across the batch, which is almost never what you meant. This is the classic softmax footgun.
• log_softmax - False by default. Flip to True for LogSoftmax, which pairs with NLLLoss or cross-entropy-based setups.

Output: one TORCH_MODEL.

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; the node is under ETK/pytorch. No model downloads - everything is built in-graph. Requirements are mostly already satisfied by ComfyUI (torch, numpy, torchvision); the file also lists scipy, scikit-learn, transformers, einops.

Common issues

• Softmax in the middle of the network. Unless you're doing something deliberate (like an attention-ish weight), a softmax buried between hidden layers destroys gradient flow. It belongs at the head.
• dim confusion. Misnormalizing over the batch instead of the classes makes every prediction sum to ~1 but the "highest" value meaningless. Normalize over the class dimension.
• Double-softmax with CrossEntropyLoss. If your loss refuses to go down, check whether you added a softmax layer before a loss that already applies one. Drop the Softmax node, or keep it only for inference by taking it off the training path.

This pack has no community tutorials to fall back on, but these nodes are thin wrappers over stock torch - the softmax rules are the same ones you'd find in any PyTorch course. And note the pack-wide quirk: it patches ComfyUI's validator to ignore return_type_mismatch, so a mis-wired connection may not give you the error message you'd expect - check wire types yourself.