Ptn KL Div Loss

KL divergence is the loss for matching distributions rather than classifying into a bucket - distillation, variational models, and any "make my predicted distribution look like the target distribution" task. PtnKLDivLoss wraps nn.KLDivLoss and, like its NLL sibling, is picky about input format in a way that bites people. The rule, straight from the author's docstring: the input (your prediction) must always be in log-probability form; the target can be either probabilities or log-probabilities depending on a flag.

How it works

KL divergence sums p(x) * log(p(x)/q(x)) over outcomes, measuring how much information you lose treating q as p. Here your prediction is q (the input) and the target is p. Because it's asymmetric - swapping the two gives a different number - you want to be deliberate about which side is which. The node passes through to nn.KLDivLoss with your chosen reduction and log_target flag.

The inputs that matter

• reduction - batchmean (sum then divide by batch size - the PyTorch-recommended choice for KL), mean, sum, or none. For KL specifically, batchmean is the mathematically proper one because plain mean scales differently than you'd expect; the docstring calls it out as recommended.
• log_target (default False) - set to True only if your target is already in log-probability form. Your prediction (input) is log-probability regardless - that part isn't optional.

Output is a single PTLOSS, wired into a trainer's loss_function or a compute-loss node. The log-probabilities must come from your model chain - there's no transformation inside this node.

Installing

Same as every node in the pack. ComfyUI Manager → search "Pt-Wrapper", or:

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

Restart ComfyUI; the pack's requirements install on first launch.

Where people get burned

Three classic traps. Feeding raw probabilities as the input instead of log-probabilities (log of a probability is negative, so this turns the loss into nonsense). Getting log_target backwards when your target is a LogSoftmax output. And reaching for mean out of habit instead of batchmean, which for KL is the one that behaves like an average per sample. Check the format before you check anything else - this loss doesn't error on wrong input, it just trains garbage quietly.