NNT Save Model

You trained a network in ComfyUI and it learned something. Now what? NNT Save Model is the node that turns your freshly-trained toy model into an actual file on disk - in seven different formats, from a plain PyTorch checkpoint to ONNX and quantized weights. It's the "export" button for the Neural Network Toolkit, and it also happens to be the cleanest way to find out which formats your hand-built nn.Module survives intact.

Like every node in this pack, the MODEL input is not a diffusion checkpoint. It's the raw PyTorch model that comes out of NNT's compile node (the thing you build from layer nodes). You can't plug a .safetensors SDXL model into this and expect it to work - wrong universe entirely.

How it works

The node takes your model, picks a save_format, and calls the matching torch serialization routine. If model_path is left empty (the default), it saves into ComfyUI/models/nnt_models/, which it registers as a proper ComfyUI model folder on first use - so it shows up in the file browser.

The formats and what they're for:

• PyTorch Model - saves the whole model plus state dict and a config blob. The most forgiving; this is what you want for resuming.
• State Dict - just the weights. You'll need the architecture saved separately to load it back.
• TorchScript - a scripted, self-contained .pt that runs without the original class definition. Requires the model to have an input_shape attribute (the compile node sets it), because it traces a random example input.
• ONNX - same input_shape requirement, exported with a dynamic batch axis, good for other runtimes.
• TorchScript Mobile - the .ptl lite-interpreter format for phones.
• Quantized - dynamic or static int8 quantization, saved as .quantized.pth. Static mode runs a quick calibration forward pass on a random input. One honest wrinkle: the quantization_bits field (4–32) is declared here but this node ignores it and always emits qint8; the bits setting actually gets used by the pack's layer-editing node instead. Don't go hunting for 16-bit output from this node.
• SafeTensors - the weights as .safetensors, with the architecture config stashed in a sidecar .config file.

There's also save_optimizer and an optimizer dropdown. Honest caveat: when you enable it, the node builds a fresh optimizer from the model's parameters at save time. It does not save the optimizer state from your actual training run, so "resume training where I left off" isn't really what this gives you. Keep that in mind before you expect mid-training checkpoints.

Outputs are MODEL (the model passes straight through, so you can keep the graph flowing) and report, a string telling you where the file landed.

Install and gotchas

Installation is the whole-pack deal:

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

Restart ComfyUI afterward, or use Manager and search "ComfyUI Neural Network Toolkit NNT". The heavy requirements (onnx, safetensors, transformers, statsmodels, pinned shap 0.41.0) mean the first pip install takes a few minutes.

The classic failure mode here: ONNX or TorchScript export dies with "Model has no input_shape attribute" when you skip the compile node or build the model in a way that never stored the input shape. If that happens, go back to your NntCompileModel and make sure it ran with a real layer stack that includes an input layer. Also, static quantization runs a calibration pass with a random tensor - fine for a toy model, a real disaster for anything where calibration data matters.