ComfyUI Node

Pt Unsqueeze

Grow a dimension where you need one

By HowToSD·Created about a year ago·Updated about a year ago· 7
Pt Unsqueeze
  • tens
  • TENSOR
dim-1

Most shape errors in tensor work aren't about big, mysterious transformations - they're about a missing axis. Pt Unsqueeze is the node that adds one: it inserts a dimension of size 1 wherever you point it, and it's the quiet fix behind half of the "shapes don't match" problems in ComfyUI-Pt-Wrapper workflows.

ComfyUI-Pt-Wrapper (HowToSD's 200-node PyTorch training pack, the spin-off of ComfyUI-Data-Analysis) talks in TENSORs, and its model nodes are strict about layout. A token sequence is (512,) but your embedding node wants a batch (1, 512). A single image from Pt From Image Transpose is (C, H, W) but the ResNet wants (N, C, H, W). Pt Unsqueeze is the one-liner that fixes exactly these. It's the node you reach for when Concat or a model node says "rank mismatch."

How it works. It's torch.unsqueeze(tens, dim): insert a size-1 axis at position dim. The source's own example: a (3, 4) tensor with dim=1 becomes (3, 1, 4).

  • tens - any TENSOR.
  • dim - where to insert the new axis. Default -1, which appends at the end: (512,)(512, 1). To prepend a batch axis you want dim=0: (3, 28, 28)(1, 3, 28, 28). Negative dims count from the end, matching PyTorch convention, and the valid range is −10 to 10 in the UI.

Output is TENSOR with one extra, size-1 axis.

The mental model that stops the confusion: unsqueeze adds a dimension of length 1 - it does not reshape, resize, or repeat anything. If you see a shape error complaining about rank (too few dimensions), unsqueeze is the fix. If it complains about sizes (a 512 vs a 768), unsqueeze will not help you - that's a different bug. And one unsqueeze gives one axis; if you need two (say (512,)(1, 1, 512)), chain the node twice or check whether the pack has a reshape node for a single-step fix.

Where you'll actually use it: every batch-prep step in the training examples. The dog-vs-cat and ResNet flows unsqueeze single images into batches; the text flows unsqueeze token sequences so the model sees a proper (batch, seq) shape. It's also the standard precursor to Concat when one side is missing its batch axis.

Install: ComfyUI Manager → "ComfyUI-Pt-Wrapper", or:

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

then restart. No model downloads; the pack's heavy requirements.txt is the only setup cost.

Troubleshooting: "dim out of range" - your tensor is smaller than the position you picked; remember dims count from 0 and negatives from the end. Size mismatch errors persist after unsqueezing - the problem was a size, not a rank; look at whether you need reshape or a different data path. Output looks identical to input when previewing - of course it does; you added an invisible axis of length 1, and it changes shape, not values. If a node still rejects it, count the dims again - you may need two unsqueezes.

CategoryData Analysis

Inputs (2)

NameTypeDefaultDescription
tensTENSOR
dimINT-1-10–10

Outputs (1)

NameTypeDescription
TENSORTENSOR