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

Pt Masked Select

Grab exactly the elements your mask says yes to

By HowToSD·Created about a year ago·Updated about a year ago· 7
Pt Masked Select
  • tens
  • masked_tens
  • TENSOR

This is the payoff node for the whole comparison-and-logic chain. Pt Masked Select runs torch.masked_select: you hand it a tensor and a boolean mask, and it returns a flat tensor containing only the elements where the mask is true. Condition says keep it, it gets kept; everything else is dropped. It's how you filter a dataset, a batch, or a feature map entirely in the graph.

How it works

Two required inputs: tens (the data) and masked_tens (the selector). The node checks that the mask is actually a boolean tensor - torch.bool - and raises a clear ValueError if you feed it floats or ints, then returns torch.masked_select(tens, masked_tens).

The docs' example is the clearest way to see it:

tens = [[1, 2, 3],
        [40, 50, 60],
        [700, 800, 900]]

masked_tens = [[True, False, False],
               [False, True, False],
               [False, False, True]]

output = [1, 50, 900]

The surprise: output is always 1D

Here's the gotcha that trips everyone on the first run: the output is flattened to one dimension. The mask's shape tells the node which elements to keep, but the result comes back as a flat vector - the original shape is gone. If you selected the diagonal of a 3×3 matrix, you don't get a diagonal matrix back; you get a 3-element vector. That's how masked_select works in torch, and it's by design here, but it means whatever you wire this into needs to expect a flat list. If you wanted to keep the grid structure, you'd be looking at a mask-multiply instead (mask times tensor, using Pt Mul).

The other requirement: the mask must be bool-typed. Good news - the output of comparison nodes like Pt Le and Pt Lt is already torch.bool, so a chain of comparison → logical AND → Pt Masked Select just works. Only hand-made 0/1 float masks need a cast first.

Installing

Pt Masked Select is part of the HowToSD/ComfyUI-Pt-Wrapper pack under the "Data Analysis" menu. Install the pack once, get ~200 nodes. ComfyUI Manager: search ComfyUI-Pt-Wrapper, install, restart. Manual:

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

The heavy requirements.txt (transformers, peft, accelerate…) serves the training side; this node needs only PyTorch, already installed. Skip the pip line for math-only work. No models to download.

Pack-wide: the TENSOR type is separate from ComfyUI's IMAGE/LATENT - convert with Pt From Image (Pt From Image Transpose for (b, c, h, w)) and Pt To Image.

CategoryData Analysis

Inputs (2)

NameTypeDefaultDescription
tensTENSOR
masked_tensTENSOR

Outputs (1)

NameTypeDescription
TENSORTENSOR