Deep Learning

torch.flatten() Explained: Reshape Tensors Like a Pro

July 9, 2026 · 23 min read
flatten in pytorch

Master torch.flatten() in PyTorch! Understand how to reshape tensors, collapse dimensions, and fix shape mismatches in your deep learning models. Learn with clear examples.

Table of Contents

Ever felt like your deep learning model’s data is a giant, messy LEGO set, and you need to pack it all neatly into a single box before feeding it to the next stage? That’s exactly what torch.flatten() does! Think of it as your ultimate data organizer in PyTorch.

You’ve probably encountered tensors with all sorts of shapes – maybe a stack of images, or the output from a convolutional layer looking like a multi-dimensional grid. But when you need to connect that to a dense layer that expects a simple, long list of numbers, you’ve got a shape mismatch problem. torch.flatten() is your hero here, taking those complex shapes and squishing them down into a manageable 1D vector, or even a tensor with fewer dimensions.

We’ll break down exactly how it works, why it’s crucial, and tackle those tricky start_dim and end_dim parameters that often leave people scratching their heads.

Real-World Examples: Where Flattening is Essential

Before we dive into the code, let’s see where torch.flatten() is a silent hero in applications you might know. This isn’t just a theoretical function; it’s a critical step in making real-world AI work.

1: Zomato/Swiggy Food Image Classification

When you upload a photo of a dish to an app like Zomato or Swiggy, a sophisticated deep learning model springs into action to identify what you’re craving. This process often begins with a Convolutional Neural Network (CNN). The CNN excels at processing the raw pixel data of an image, progressively extracting features like edges, textures, and shapes.

The output of these early convolutional layers isn’t a simple list of numbers; it’s a collection of multi-dimensional “feature maps” that represent these detected patterns. To make a final decision about the dish’s identity (e.g., “Paneer Butter Masala” versus “Masala Dosa”), these complex, multi-dimensional feature maps must be flattened. This means collapsing all the spatial and feature dimensions into a single, long vector.

This flattened vector then serves as the input to the final classification layer (a fully connected or linear layer) of the neural network, which uses this comprehensive feature representation to predict the most likely food item. Without torch.flatten(), the output of the CNN would be incompatible with the subsequent linear classification layer.

2: Indian Monsoon Rainfall Prediction (IMD Data)

The Indian Meteorological Department (IMD) uses complex models to forecast monsoon rainfall, a critical aspect of India’s economy and agriculture. These models ingest vast amounts of data, including historical weather patterns, satellite imagery, and readings from meteorological stations.

This data, when processed, can be represented as multi-dimensional tensors. For instance, a tensor might capture information over time, across geographical grids (latitude and longitude), and include various atmospheric features (temperature, pressure, humidity). After these raw data points are processed by initial layers that might extract temporal trends or spatial correlations, the resulting multi-dimensional structure often needs to be flattened.

This flattening transforms the grid-like or time-series data into a single, manageable feature vector. This vector then feeds into a predictive model (often a form of regression or classification) that outputs a specific forecast, such as the expected rainfall in millimeters for a particular region or the probability of above-average rainfall. torch.flatten() is essential here to bridge the gap between the spatially and temporally rich input data representation and the single-value output required for the final prediction.

3: Duolingo’s Adaptive Learning Models

Duolingo’s success lies in its ability to personalize the learning experience for millions of users. At the heart of this personalization are adaptive learning models. When you answer questions, Duolingo’s system collects a sequence of data points about your performance: whether you got an answer right or wrong, the type of question, how long you took, and perhaps even your previous performance on similar items.

This rich, contextual information is often represented as a complex, multi-dimensional tensor or a sequence of vectors. To determine the optimal next exercise for you, this input data needs to be summarized into a concise representation of your current learning state. This is where flattening comes in.

The model processes the input sequence and then flattens it into a single feature vector. This vector encapsulates your learned knowledge, your current difficulties, and your learning trajectory. This flattened vector is then fed into a decision-making component of the model, which uses it to select the most appropriate exercise from Duolingo’s vast library, ensuring that the learning path is challenging yet achievable for you.

The Core Concept: From a Grid to a Line

At its heart, torch.flatten() is a reshaping tool. Its main job is to take a multi-dimensional tensor and collapse some or all of its dimensions into a single one.

Imagine you have a batch of 10 grayscale images, each 28×28 pixels. Your data’s shape would be (10, 28, 28). A standard neural network layer (a “Linear” or “Fully Connected” layer) can’t handle this 3D grid. It expects a 2D tensor where each row is a single, long list of features for one sample. flatten() is the bridge. It transforms (10, 28, 28) into (10, 784), where each of the 10 images is now represented by a single vector of 784 pixels (28×28).

Key Term 1: torch.flatten()

This is the main tool for the job. Think of it as a steamroller for your data’s dimensions. It takes a bulky, multi-dimensional tensor and flattens it.

WHAT is torch.flatten()?

torch.flatten() is a function in PyTorch that transforms a tensor by collapsing a specified range of its dimensions into a single new dimension. For example, it can take a tensor representing a batch of color images with a shape like (32, 3, 28, 28) and “squish” the channel, height, and width dimensions together. The result would be a 2D tensor of shape (32, 2352), where each of the 32 images is now represented by a single, long list of 2,352 numbers (3×28×28).

WHY does it exist?

Its primary purpose is to act as a bridge between different types of neural network layers.

  1. Shape Compatibility: Convolutional layers (like nn.Conv2d) are great at finding patterns in grids, like images. They output multi-dimensional tensors called feature maps. However, fully connected layers (like nn.Linear) expect their input for each sample to be a simple 1D list (a vector). torch.flatten() solves this shape mismatch by converting the grid-like output of the convolutional layers into the list-like input that the linear layers need.

  2. Feature Vector Creation: It effectively creates a single “feature vector” from all the information extracted by earlier layers. This vector summarizes the entire input (e.g., an image) in a format that’s easy for the final classification or regression part of the model to use.

HOW does it work?

Mechanically, torch.flatten() reads the elements from the original tensor in a specific order (contiguously in memory) and rearranges them into the new, flattened shape. It starts from the first dimension to be flattened (start_dim) and proceeds through the last (end_dim).

Importantly, torch.flatten() is highly efficient. It usually returns a “view” of the original tensor, not a copy. This means it doesn’t create a new block of memory for the data; it just creates a new “header” that describes how to interpret the existing data in a different shape. It only makes a copy if the original tensor’s data is not laid out in a continuous block in memory, ensuring the operation always succeeds.

WHEN to use it (and when not to)?

WHERE is it applied in practice?

You will find torch.flatten() in two forms:

  1. As a function: torch.flatten(x, start_dim=1) is called directly within the forward method of a model. This is a flexible, on-the-fly approach.

  2. As a layer: nn.Flatten(start_dim=1) is defined as a layer in your model’s __init__ method. This is often considered cleaner and is the preferred way when building models using nn.Sequential, as it makes the reshaping step an explicit part of the model’s architecture.

ADVANTAGES

LIMITATIONS

Key Term 2: start_dim (parameter)

This parameter is your control knob for telling flatten where to start its work. Getting this right is the key to using the function correctly.

WHAT is start_dim?

start_dim is an integer parameter for the torch.flatten() function that specifies the first dimension in the flattening range. Dimensions are indexed starting from 0. Any dimension with an index before start_dim will be left untouched. The flattening process begins at the dimension specified by start_dim.

WHY does it exist?

It exists to provide crucial control, primarily to preserve the batch dimension. In deep learning, we almost always process data in batches (e.g., 32 images at a time). The shape of our tensor is typically (batch_size, channels, height, width). We want to flatten each image individually but keep them separate within the batch. By setting start_dim=1, we tell PyTorch: “Leave dimension 0 (the batch dimension) alone, and start flattening from dimension 1 (the channels dimension) onwards.”

HOW does it work?

The function looks at the start_dim value you provide. Let’s say your tensor has a shape (A, B, C, D) and you call flatten(tensor, start_dim=1).

  1. The function sees start_dim=1.

  2. It identifies dimension 0 (with size A) as being before the start index. It decides to keep this dimension as is.

  3. It then takes all dimensions from the start_dim (index 1) to the end_dim (by default, the last one) and multiplies their sizes together. In this case, it combines dimensions B, C, and D.

  4. The final output shape becomes (A, B \times C \times D).

WHEN to use it (and when not to)?

WHERE is it applied in practice?

It is used as the second argument to the torch.flatten function or the first argument to the nn.Flatten layer constructor.

ADVANTAGES

LIMITATIONS

Key Term 3: end_dim (parameter)

This is the other control knob. It tells flatten where to stop its work, allowing for more surgical reshaping.

WHAT is end_dim?

end_dim is an integer parameter that specifies the last dimension to be included in the flattening operation. The function will collapse all dimensions from start_dim up to and including end_dim. Dimensions that come after end_dim will be left untouched. The default value is -1, which is a special index in Python that means “the very last dimension.”

WHY does it exist?

It provides fine-grained control for situations where you don’t want to flatten everything to the end of the tensor. For example, you might have a tensor representing a batch of sentences, where each word has an embedding. The shape could be (batch, num_sentences, words_per_sentence, embedding_dim). You might want to flatten the words_per_sentence and embedding_dim together to get a “sentence vector” but keep the num_sentences dimension separate. end_dim allows you to do this.

HOW does it work?

When you call flatten(tensor, start_dim, end_dim), the function identifies the range of dimensions from the start_dim index to the end_dim index. It then multiplies the sizes of all dimensions within this range to create a single new dimension. The dimensions before start_dim and after end_dim remain in their original state.

For a tensor of shape (A, B, C, D, E) and a call flatten(tensor, start_dim=1, end_dim=2):

  1. Dimension 0 (size A) is preserved.

  2. Dimensions 1 and 2 (sizes B and C) are flattened together.

  3. Dimensions 3 and 4 (sizes D and E) are preserved.

  4. The final output shape will be (A, B \times C, D, E).

WHEN to use it (and when not to)?

WHERE is it applied in practice?

It is used as the third argument to the torch.flatten function or the second argument to the nn.Flatten layer constructor.

ADVANTAGES

LIMITATIONS

Key Parameters at a Glance

When you call torch.flatten(), these are the only parameters you need to know:

Example Scenarios:

PyTorch Code in Action

Let’s see how this works with a hands-on example.

python
import torch
import torch.nn as nn

# --- Basic Flattening ---

# 1. Create a sample tensor representing a batch of 4 images (3 channels, 2x2 pixels)
# Shape: (batch_size, channels, height, width)
input_tensor = torch.randn(4, 3, 2, 2)
print(f"Original Tensor Shape: {input_tensor.shape}")

# 2. Flatten the entire tensor (default behavior)
# This is usually INCORRECT for batches
flat_all = torch.flatten(input_tensor)
print(f"\nShape after flatten(input_tensor): {flat_all.shape}")  # Flattens everything

# 3. Flatten correctly, preserving the batch dimension (start_dim=1)
# This is the MOST COMMON use case
flat_correct = torch.flatten(input_tensor, start_dim=1)
print(f"Shape after flatten(input_tensor, start_dim=1): {flat_correct.shape}")

# 4. Flatten a specific range of dimensions (e.g., channels and height)
# Shape: (batch, channel, height, width) -> (batch, channel*height, width)
flat_partial = torch.flatten(input_tensor, start_dim=1, end_dim=2)
print(f"Shape after flatten(input_tensor, start_dim=1, end_dim=2): {flat_partial.shape}")


# --- Using nn.Flatten in a Model ---

# A simple model that mimics going from a convolution to a linear layer
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        # Imagine this is our convolutional layer's output
        self.conv_output_shape = (1, 64, 8, 8)  # (Batch, Channels, H, W)

        # Use nn.Flatten to handle the reshaping automatically
        # We tell it to start flattening from dimension 1
        self.flatten_layer = nn.Flatten(start_dim=1)

        # The linear layer now needs to accept the flattened size
        # 64 * 8 * 8 = 4096
        self.linear_layer = nn.Linear(in_features=4096, out_features=10)  # 10 output classes

    def forward(self, x):
        print(f"Shape before flatten: {x.shape}")
        x = self.flatten_layer(x)
        print(f"Shape after flatten: {x.shape}")
        x = self.linear_layer(x)
        print(f"Shape after linear layer: {x.shape}")
        return x

# Let's test the model
model = SimpleCNN()
# Create a fake output from a conv layer
fake_conv_output = torch.randn(1, 64, 8, 8)
model(fake_conv_output)

Common Mistakes and How to Avoid Them

1. Mistake: Flattening the Batch Dimension

A very common error is calling torch.flatten(x) with its default start_dim=0. If your input x has a shape like (32, 3, 28, 28) (a batch of 32 images), this will collapse everything into one giant vector, mixing all 32 images together.

Fix: Always specify start_dim=1 when working with batches: torch.flatten(x, start_dim=1). This preserves the batch dimension, resulting in the correct shape (32, 2352).

❌ Before:

python
import torch

# A batch of 4 images, each with 3 channels and 2x2 pixels
batched_data = torch.randn(4, 3, 2, 2)
print(f"Original shape: {batched_data.shape}")

# MISTAKE: Using default start_dim=0 flattens the batch dimension too!
# All 4 images are mixed together into one vector.
flattened_wrong = torch.flatten(batched_data) 

print(f"Shape after WRONG flatten: {flattened_wrong.shape}")
# Expected output: Shape after WRONG flatten: torch.Size([48])

✅ After:

python
import torch

# A batch of 4 images, each with 3 channels and 2x2 pixels
batched_data = torch.randn(4, 3, 2, 2)
print(f"Original shape: {batched_data.shape}")

# FIX: Specify start_dim=1 to preserve the batch dimension.
# Each of the 4 images is flattened individually.
flattened_correct = torch.flatten(batched_data, start_dim=1)

print(f"Shape after CORRECT flatten: {flattened_correct.shape}")
# Expected output: Shape after CORRECT flatten: torch.Size([4, 12])

2. Mistake: Confusing flatten() with reshape() or view()

While x.reshape(batch_size, -1) can achieve a similar result, flatten() is more explicit and often safer. view() can fail on non-contiguous tensors, while flatten() will always work (it may return a copy if needed).

Fix: Prefer using torch.flatten(x, start_dim=1) or the nn.Flatten() layer in your models. It clearly states your intention: to flatten dimensions for a linear layer.

3. Mistake: Does it Copy Data?

Students often worry about memory. torch.flatten() returns a new tensor that often shares the same underlying data as the original (it’s a “view”). It only creates a copy if the original tensor is not contiguous in memory. You generally don’t need to worry about performance issues.

Putting It All Together: Full Code Examples

Example 1: nn.Flatten in a CNN Training Loop

This example shows a complete, practical application. We define a simple CNN where nn.Flatten is a crucial layer that connects the convolutional feature extractor to the final linear classifier. The code includes a full training step to demonstrate how it works in a real model.

python
import torch
import torch.nn as nn
import torch.optim as optim

# PyTorch Version Note: This code is written for PyTorch 2.0.0+
# For reproducibility, we set a fixed random seed.
torch.manual_seed(42)

# --- 1. Setup Device ---
# Set device to GPU ('cuda') if available, otherwise use CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
print(f"PyTorch Version: {torch.__version__}")

# --- 2. Model Definition with nn.Flatten ---
# This model simulates a common computer vision pipeline: Conv -> Flatten -> Linear
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        # A simple convolutional layer that takes in 3-channel images (e.g., RGB)
        # and outputs 16 feature maps.
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()

        # The KEY component: nn.Flatten()
        # It is configured with start_dim=1 to preserve the batch dimension (dim 0)
        # and flatten all subsequent dimensions (channels, height, width).
        self.flatten = nn.Flatten(start_dim=1)

        # A linear layer for classification.
        # The in_features must match the size of the flattened feature vector.
        # For a 32x32 input image, after conv1, the shape is (batch, 16, 32, 32).
        # The flattened size is 16 * 32 * 32 = 16384.
        self.fc1 = nn.Linear(in_features=16 * 32 * 32, out_features=num_classes)

    def forward(self, x):
        # x shape: [batch_size, 3, 32, 32]
        print(f"Initial input shape: {x.shape}")

        x = self.conv1(x)
        x = self.relu(x)
        # x shape after conv/relu: [batch_size, 16, 32, 32]
        print(f"Shape after convolution: {x.shape}")

        # Flatten the feature maps
        x = self.flatten(x)
        # x shape after flatten: [batch_size, 16384]
        print(f"Shape after flatten: {x.shape}")

        # Pass through the final linear layer
        x = self.fc1(x)
        # x shape after linear layer: [batch_size, num_classes]
        print(f"Final output shape: {x.shape}\n")
        return x

# --- 3. Data, Model, Loss, and Optimizer Initialization ---
# Create a dummy batch of data
batch_size = 4
# Dummy images (4 images, 3 channels, 32x32 pixels)
dummy_input = torch.randn(batch_size, 3, 32, 32).to(device)
# Dummy labels (4 labels, for 10 classes)
dummy_labels = torch.randint(0, 10, (batch_size,)).to(device)

# Instantiate the model and move it to the selected device
model = SimpleCNN(num_classes=10).to(device)

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

# --- 4. Complete Training Loop ---
# This demonstrates a single training step where flatten is used.
print("--- Running a single training step ---")

# a. Zero the gradients
# Always call this before the backward pass to clear old gradients
optimizer.zero_grad()

# b. Forward pass
# Get model predictions. The forward method will print the shape changes.
outputs = model(dummy_input)  # Shape: [4, 10]

# c. Calculate loss
loss = criterion(outputs, dummy_labels)
print(f"Calculated Loss: {loss.item():.4f}")

# d. Backward pass
# Compute gradients of the loss with respect to model parameters
loss.backward()

# e. Optimizer step
# Update the model's weights based on the computed gradients
optimizer.step()

print("\n--- Training step complete ---")
print("Model weights have been updated.")

Example 2: Exploring start_dim and end_dim

This example isolates the torch.flatten function itself to demonstrate how the start_dim and end_dim parameters provide fine-grained control over the reshaping process. This is useful for understanding the function’s mechanics before using it inside a larger model.

python
import torch

# PyTorch Version Note: This code is written for PyTorch 2.0.0+
# For reproducibility, we set a fixed random seed.
torch.manual_seed(42)

# --- Setup Device ---
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
print(f"PyTorch Version: {torch.__version__}\n")

# --- Create a sample tensor ---
# Shape represents: (batch_size, channels, height, width)
# Let's imagine a batch of 2 images, each with 3 channels and 4x5 pixels.
input_tensor = torch.arange(120).reshape(2, 3, 4, 5).to(device)
print(f"Original Tensor Shape: {input_tensor.shape}")  # Shape: [2, 3, 4, 5]

# --- Scenario 1: Default Flatten (Incorrect for Batches) ---
# Flattens the ENTIRE tensor into a single 1D vector.
# start_dim=0 and end_dim=-1 are the defaults.
flat_all = torch.flatten(input_tensor)
# The batch dimension is lost. 2 * 3 * 4 * 5 = 120
print(f"\n1. Shape after default flatten(input): {flat_all.shape}")  # Shape: [120]

# --- Scenario 2: Standard Use Case (Correct for Batches) ---
# Flattens all dimensions starting from dimension 1.
# This preserves the batch dimension (dim 0).
flat_correct = torch.flatten(input_tensor, start_dim=1)
# The new shape is (batch_size, features). 3 * 4 * 5 = 60
print(f"2. Shape after flatten(input, start_dim=1): {flat_correct.shape}")  # Shape: [2, 60]

# --- Scenario 3: Partial Flatten with end_dim ---
# Flattens only a specific range of dimensions.
# Here, we flatten from dim 1 (channels) to dim 2 (height).
# The width dimension (dim 3) is left untouched.
flat_partial = torch.flatten(input_tensor, start_dim=1, end_dim=2)
# The new shape combines channels and height: 3 * 4 = 12
print(f"3. Shape after flatten(input, start_dim=1, end_dim=2): {flat_partial.shape}")  # Shape: [2, 12, 5]

# --- Scenario 4: Flattening only the batch and channel dimensions ---
# Flattens from dim 0 (batch) to dim 1 (channels).
# The spatial dimensions (height, width) are preserved.
flat_batch_channel = torch.flatten(input_tensor, start_dim=0, end_dim=1)
# The new shape combines batch and channels: 2 * 3 = 6
print(f"4. Shape after flatten(input, start_dim=0, end_dim=1): {flat_batch_channel.shape}")  # Shape: [6, 4, 5]

Summary and Key Takeaways

0 0 votes
Article Rating
Subscribe
Notify of
guest
0 Comments
Oldest
Newest Most Voted
// STILL BROWSING?
Build along, don't just read.
Get labs & articles matched to what you're into — free, takes 30 seconds.
Start building free
0
Would love your thoughts, please comment.x
()
x