FlattenTensor
The one-node trick that keeps shape mismatches out of your models
- 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 downstreamnn.Linear. For a batch you usually want shape-agnostic flattening inside the model instead (AddReshapeLayerwith1, -1keeps the batch axis meaningful in a(1, -1)sense). - Inference vs training shapes. Flattening before a linear layer means the linear's
in_featuresmust 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 isTORCH_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.
Inputs (1)
| Name | Type | Default | Description |
|---|---|---|---|
| input_data | TORCH_TENSOR | — |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| TORCH_TENSOR | TORCH_TENSOR | — |