Nodes/ComfyUI-Pt-Wrapper/Pt Index Select
ComfyUI Node

Pt Index Select

Pull the exact rows you want out of a tensor

By HowToSD·Created about a year ago·Updated about a year ago· 7
Pt Index Select
  • tens
  • TENSOR
dim0
index

If you're poking around in the pack's model-building workflows, this is the node you'll reach for whenever you have a tensor full of stuff and only want some of it. Say you've got a batch of embeddings, one row per image or per class, and you want to keep rows 0 and 2 but drop the rest. Pt Index Select is torch.index_select wearing a ComfyUI jacket: it grabs the slices you name, along the axis you name, and hands back a tensor that keeps its shape.

How it works

Under the hood it's dead simple. The index box holds a string that looks like a Python list - [0, 2] - which the node parses and turns into an int64 index tensor, then runs torch.index_select(tens, dim=dim, index=index_tensor). That's the whole mechanism. The output keeps the same rank as the input; indexing never squeezes anything away. The pack's docs walk through a concrete example: with a 3×3 tensor and index = "[0, 2]", dim = 0, you get back the first and third rows, each still a full row.

The inputs that matter

There are only three, and two of them are where beginners stumble:

  • tens - the TENSOR you're slicing.
  • dim - which axis to slice along, 0-indexed. Default 0. For a 2D tensor, dim=0 selects rows, dim=1 selects columns.
  • index - the selector. This is the trap: it must be a list literal, brackets included. [0, 2] works; 0, 2 does not, because the parser expects valid Python (ast.literal_eval). It's multiline, so a long index list is fine to paste in.

The output is a single TENSOR, ready to wire into any other Pt node.

Gotchas

Indices are 0-based, so the first row is index 0, not 1. And if an index is out of range for the size of that dimension, index_select throws - you'll get a red node rather than a silent wrong answer, which is honestly the better failure mode.

One distinction worth making: Pt Index Select picks by position, where Pt Masked Select picks by a boolean mask and flattens the result into 1D. If you need conditional logic ("keep everything above a threshold") rather than a fixed list, you want the masked version. If you want a fixed subset and its shape preserved, this is the one.

Installing

This ships inside the HowToSD/ComfyUI-Pt-Wrapper pack (it lands under the "Data Analysis" menu), so installing the pack gets you all ~200 of its nodes at once. Easiest route: ComfyUI Manager → "Install Custom Nodes" → search ComfyUI-Pt-Wrapper → Install, then restart.

Manual route:

cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper
cd ComfyUI-Pt-Wrapper
pip install -r requirements.txt

Heads up: that requirements list is heavy (transformers, datasets, peft, accelerate, scikit-learn…). This node only needs PyTorch, which ComfyUI already has - the heavy stuff is for the pack's training side. If you're just doing tensor math you can skip the pip line and restart. No model downloads either way.

Also know the pack's shared gotcha: it speaks a custom TENSOR type, not ComfyUI's IMAGE or LATENT. To get image data in, convert with Pt From Image (or Pt From Image Transpose for the (b, c, h, w) layout), and Pt To Image to get back out. Wires won't connect across types until you do.

CategoryData Analysis

Inputs (3)

NameTypeDefaultDescription
tensTENSOR
dimINT0-10–10
indexSTRING

Outputs (1)

NameTypeDescription
TENSORTENSOR