Ptn Model With Closure

Sometimes a model's raw output isn't what you want - you want a sigmoid after the logits, a softmax over classes, a log-softmax for an NLL loss. PtnModelWithClosure is the node that glues a function onto the end of a model's forward pass: you give it a model and a PTCALLABLE (an object that wraps a function), and it returns a new PTMODEL that runs the callable on the model's output. It's how you customize the tail of a model without writing a custom layer.

How it works

It builds a wrapper module holding your model and a closure. On every forward pass it runs x = model(x), then x = closure(x), and returns the result. The closure is fully differentiable and part of the graph, so gradients flow through it during training. The key distinction: a closure is a stateless function applied to the output, not a layer with weights - so this is the right tool for activations (which have no parameters), not for learned post-processing.

The inputs

• model - any PTMODEL.
• closure - a PTCALLABLE, produced by the pack's activation nodes: Ptf Sigmoid, Ptf Softmax, Ptf Log Softmax, Ptf ReLU, and friends all emit PTCALLABLE outputs.

One PTMODEL comes out, and it's a drop-in replacement for the model you put in - chain it, train it, save it, exactly as if it were the original.

How you'd use it

The clearest use is pairing with PtnNLLLoss: that loss demands log-probabilities, so wrap your model with Ptf Log Softmax as the closure and the output is ready to feed it. Or, if a workflow expects probabilities (0–1) out of a model whose final layer is a bare nn.Linear, wrap it with Ptf Sigmoid. The pack's loss nodes like Ptn BCE With Logits Loss already absorb the sigmoid internally - but when your model chain itself needs to emit the transformed values (for downstream nodes, not the loss), this is the piece that does it.

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

Feeding a PTMODEL where a PTCALLABLE goes (or vice versa) is the one real mistake, and ComfyUI will refuse the wire rather than silently cope. Also remember the closure runs unconditionally - if you wrap a model with a sigmoid and later connect it to Ptn BCE With Logits Loss (which sigmoids internally), you've applied it twice and your training will crawl. Pick one place for the activation.