Skip to content

Flatten

Bases: BaseModule

Flatten a 4-D spatial tensor into a 2-D flat tensor.

The tensor is permuted from [B, C, H, W] to [B, H, W, C] before being flattened along the trailing three dimensions (for compiler compatibility).