ComfyUI Node

FlattenTensor

The one-node trick that keeps shape mismatches out of your models

By TashaSkyUp·Created about a year ago·Updated about a year ago· 1
FlattenTensor
  • input_data
  • TORCH_TENSOR

Fully-connected layers only accept one-dimensional feature vectors per sample. Your images are (batch, height, width, channels). Somewhere between them, the tensor has to get flattened - and FlattenTensor is the pack's one-node way to do it. Feed it any TORCH_TENSOR and it comes out as a single 1D vector (the whole thing, batch included, flattened to (N,)). It's the tensor equivalent of the flatten step every CNN tutorial writes before the first nn.Linear.

You reach for it right before data hits a model: take a ComfyUIImageToPytorchTENSOR output, flatten it, and feed the result to a training or inference node - or flatten features to hand them to TrainModel. It's the most boring node in the pack and also one of the most-used, because it eliminates a whole class of "expected 2D input, got 4D" errors.

How it works

It's literally input_data.flatten() - PyTorch's all-dims flatten, which turns (2, 3, 4) into (24,). No options, no start/end dims, no debate. If you need a partial flatten (keep the batch, flatten the rest) that's not what this node does - you'd use a reshape elsewhere in the pack, or AddReshapeLayer inside the model.

Input and output

  • input_data (TORCH_TENSOR) - anything tensor-shaped: images, dataset features, model outputs.
  • Output: one TORCH_TENSOR, fully flattened to a single dimension.

Install

ComfyUI Manager, search "EternalKernel PyTorch Nodes", or:

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

Restart ComfyUI; node under ETK/pytorch. No model files. Requirements are the standard stack plus scipy, scikit-learn, transformers, einops.

Common issues

  • It flattens the batch too. If your data is (64, 784) and you were hoping for per-sample flattening, the output is (50176,) - one giant vector, which changes what "sample" means to a downstream nn.Linear. For a batch you usually want shape-agnostic flattening inside the model instead (AddReshapeLayer with 1, -1 keeps the batch axis meaningful in a (1, -1) sense).
  • Inference vs training shapes. Flattening before a linear layer means the linear's in_features must match the flattened size exactly. Change the input resolution and you break it - count dimensions.
  • Wire-type gotcha. The pack patches ComfyUI's validator to ignore return_type_mismatch, so a tensor wired into a node expecting something else may not raise the error you'd expect. FlattenTensor output is TORCH_TENSOR; check the target node's input type.

There are no community tutorials for this pack - it's a quiet corner - but a flat tensor is a flat tensor, and the PyTorch docs on flatten() cover every edge case. When in doubt, count your dimensions before and after.

CategoryETK/pytorch

Inputs (1)

NameTypeDefaultDescription
input_dataTORCH_TENSOR

Outputs (1)

NameTypeDescription
TORCH_TENSORTORCH_TENSOR