AddBatchNormLayer

A deeper network in the EternalKernel PyTorch Nodes pack trains painfully slowly, or the loss stalls entirely? That's the classic sign of internal covariate shift - layer inputs drifting as weights update - and batch normalization is the standard fix. AddBatchNormLayer appends torch.nn.BatchNorm1d or BatchNorm2d to your nn.Sequential, and it works like a normalizer for the middle of your network: it rescales each layer's activations to a stable mean and variance during training, learns per-feature scale/shift parameters, and keeps those stats frozen at inference.

It's the reason modern networks can stack 50+ layers when old ones choked at 10. For the small models this pack builds, you'll reach for it when a network with more than a couple of layers stops converging, or when you want to use higher learning rates without the loss exploding.

How it works

The node builds nn.BatchNorm1d(num_features) by default, or BatchNorm2d when you flip one_d off, and inserts it at the end of your Sequential (mutating in place, returning the model). BatchNorm tracks running mean/variance during training and uses them at eval time - so a freshly-built model (before any training) has uninitialized running stats, and inference on an untrained model can behave oddly. That's PyTorch behavior, not a bug here.

Inputs

• model (TORCH_MODEL) - the Sequential to extend. Start from SequentialModelProvider.
• num_features - the number of channels (for 2D) or features (for 1D) the layer before it outputs. This is the input you'll get wrong. It must exactly match the channel count of the tensor arriving - one less and the whole graph errors on shape.
• one_d - True gives BatchNorm1d (for flat feature vectors, i.e. after linear layers); False gives BatchNorm2d (for conv feature maps). The default is True, so if you're using this after a Conv2d, flip it to False or you'll get a 1D layer fed 4D data.

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; node is under ETK/pytorch. No model files. Requirements are ComfyUI-standard (torch, numpy, torchvision) plus scipy, scikit-learn, transformers, einops.

Common issues

• num_features mismatch. BatchNorm validates against the incoming tensor's channel count and will throw a clear shape error. Match num_features to the previous layer's out_channels/out_features.
• BatchNorm with batch size 1. The TrainModel node's default batch_size is 1, and BatchNorm statistics from a batch of one are degenerate - it can train erratically or error. Bump the batch size up (16–64) when you use BatchNorm, which means building your data with DatasetToDataloader/batched tensors rather than single samples.
• 1D vs 2D confusion. BatchNorm1d on a conv tensor (4D) errors immediately. Match one_d to the layer type before it.

Small pack, no tutorials, no community footprint - but the mechanics are pure PyTorch, so the standard BatchNorm wisdom all applies. One pack quirk to remember: it patches ComfyUI's validator to ignore return_type_mismatch, so a mis-wired node may not flag itself the way you'd hope.