Hello! Today we are covering in-place operations in PyTorch.

When building deep learning models, memory management becomes a major challenge. For many of you, the biggest part of that challenge is probably GPU memory.

This is where PyTorch's in-place operations come to the rescue.

In this article, we will look at how to use in-place operations from a variety of angles.

What Are In-Place Operations?

The Basic Idea

An in-place operation is one that directly overwrites an existing memory region. In PyTorch, you can perform an in-place operation by appending an underscore (_) to the operator name.

In other words, while a regular operation needs to allocate new memory, an in-place operation can rewrite the existing memory directly.

Let's see it in action.

import torch

# Regular operation
x = torch.tensor([1, 2, 3])
y = x + 5  # New memory is required

# With an in-place operation, do this!
x = torch.tensor([1, 2, 3])
x.add_(5)  # Rewrites x directly

Now let's take a closer look at how memory is used in the code above and how the two operations differ.

First, consider the regular operation (x + 5). Here, after a tensor x holding the data [1, 2, 3] is created in memory, executing x + 5 allocates a new memory region.

The computed result [6, 7, 8] is written into this new region, and a reference to that new memory region is assigned to y.

At this point, the original x is unchanged and remains [1, 2, 3]. As a result, two distinct tensors, x and y, exist in memory.

With the in-place operation (x.add_(5)), the processing differs significantly. A tensor x holding [1, 2, 3] is created in memory, but when the add_ method is called, the computation is performed directly in the same memory region used by x.
The original data [1, 2, 3] is overwritten with [6, 7, 8], and no new memory region is allocated at all. x simply comes to point at the updated values.

Checking the actual memory addresses, as shown below, makes the difference clear

# Regular operation
x = torch.tensor([1, 2, 3])
print(f"Memory address of x: {x.data_ptr()}")
y = x + 5
print(f"Memory address of y: {y.data_ptr()}")  # A different address from x

# In-place operation
x = torch.tensor([1, 2, 3])
print(f"Address of x before the operation: {x.data_ptr()}")
x.add_(5)
print(f"Address of x after the operation: {x.data_ptr()}")  # The address is the same

The importance of this difference becomes especially pronounced when working with large tensors.

For example, consider operating on a 1000×1000 matrix

# For a 1000x1000 matrix (float32)
large_x = torch.randn(1000, 1000)  # About 4MB

With regular operations, every computation on this 4MB of data requires another 4MB of memory. With in-place operations, no additional memory is needed—the existing 4MB region is reused.

This difference in memory efficiency matters especially when running inference on large neural networks or executing in memory-constrained environments.

Next, let's look at more practical use cases.

Practical Usage

1. Optimizing Inference

The following code is a simple yet efficient network implementation

class EfficientNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)
        self.bn = nn.BatchNorm2d(64)
        
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x.relu_()  # Apply ReLU in place!
        return x

Let's look at the key points of this code.

First, regarding the model's structure: it takes 3 channels as input (assuming RGB color images) and passes them through a 3×3 convolutional layer with 64 filters.

After that, batch normalization is applied, and finally a ReLU activation function.

Pay particular attention to how ReLU is applied.

Where you would normally write x = torch.relu(x) or F.relu(x), we use the in-place operation x.relu_().

This has a significant advantage

With a regular ReLU operation, the following memory activity occurs

1. Memory for the convolutional layer's output
2. Memory for the batch normalization output
3. New memory for the ReLU

But with an in-place operation, no new memory is needed for the ReLU.

That is because it directly overwrites the batch normalization output.

During inference in particular, there is no need to keep the batch normalization output around, which makes this optimization possible.

When dealing with large tensors, as in image processing, this difference actually becomes quite significant.

For example, when processing 224×224 RGB images with a batch size of 32 in models like VGG or ResNet, the output of this layer alone saves:

32 (batch size) * 64 (channels) * 222 * 222 (output size) * 4 (bytes per float32)

that much memory.

Note that caution is required if you use this code during training.

If the ReLU gradient is needed for backpropagation, it is safer to avoid the in-place operation.

In that case, rewrite it as follows.

def forward(self, x):
    x = self.conv(x)
    x = self.bn(x)
    x = F.relu(x)  # Use the regular ReLU during training
    return x

2. Streamlining Data Preprocessing

One of the most common operations in data preprocessing is normalization. Especially when handling large datasets, the memory efficiency of this step becomes extremely important.

The following code is an example of efficient normalization using in-place operations

def efficient_normalize(tensor):
    mean = tensor.mean()
    std = tensor.std()
    tensor.sub_(mean).div_(std)  # They can be chained!
    return tensor

This simple implementation contains a few nice touches. First, mean() and std() return scalar values, so their memory usage is essentially negligible.
The important part here is that the normalization computation is done with in-place operations.

A computation that would typically be written like this,

normalized = (tensor - mean) / std  # Requires a new memory region

avoids the extra memory allocation by using in-place operations.

Furthermore, chaining sub_() and div_() keeps the code neat and compact.

For example, when normalizing image data with a batch size of 32 (224×224 RGB images),

# Batch size * channels * height * width * 4 bytes
32 * 3 * 224 * 224 * 4 = about 19MB

that is roughly how much memory you save

If normalization runs on every load in your data loader, this difference adds up quickly.

Of course, if you want to keep the original data, use it like this.

normalized_data = efficient_normalize(original_data.clone())

3. Using In-Place Operations in the Optimization Step During Training

The following code is an example implementation of SGD (stochastic gradient descent) using in-place operations. Weight updates are used all the time, so this is worth knowing.

def custom_sgd_update(parameters, lr):
    with torch.no_grad():
        for param in parameters:
            if param.grad is not None:
                param.sub_(lr * param.grad)  # Update weights in place

By using the with torch.no_grad(): context manager, we disable autograd tracking.

Where is the in-place operation?

Right—the weight update uses the in-place operation (sub_).

If we used regular subtraction:

param = param - lr * param.grad  # Requires a new memory region

a new memory region would be required for every update.

By using the in-place operation instead,

1. No new memory allocation is needed
2. Memory release and reallocation can be skipped
3. Even large models can be updated efficiently

For example, for a model with 50 million parameters, you no longer need to allocate and free a memory region of the following size on every update.

# Number of parameters * 4 bytes (for float32)
50,000,000 * 4 = about 190MB

It is used like this

# Update the model's parameters
learning_rate = 0.01
custom_sgd_update(model.parameters(), learning_rate)

Points to Watch Out For

Compatibility with Autograd

Let's look at the relationship between PyTorch's autograd system and in-place operations through concrete code

x = torch.tensor([1., 2., 3.], requires_grad=True)
y = x * 2

# This is fine!
z = y.relu()

# This is dangerous!
y.relu_()  # The information needed for gradient computation gets lost...

Let's use this code to examine the pitfalls of in-place operations in detail. First, we create a tensor x and enable gradient computation by specifying requires_grad=True. Next, we multiply x by 2 and store the result in y.

The important point here is that PyTorch's autograd system is building a computation graph. When using the regular ReLU function (z = y.relu()), a new tensor z is created and the original value y is retained in the computation graph. This allows gradients to be computed correctly during backpropagation.

However, using the in-place relu_() overwrites the original value y directly. In this case, the information needed at backpropagation time (the values before the activation function was applied) is lost, so correct gradient computation becomes impossible.

Specifically, in the ReLU gradient computation, the information needed to determine whether the input was positive or negative disappears.

In practice, PyTorch may raise a runtime error for such dangerous operations. This is an important safety mechanism to protect the consistency of the computation graph. Intuitively, a regular operation that produces a new result while preserving the original value is a "history-preserving" operation, while an in-place operation is a "history-overwriting" one.

So Is Using In-Place Operations During Training Dangerous or Not?

It is fine to use regular operations in the forward pass during training and to restrict in-place operations to situations that do not affect gradient computation—such as inference, or the weight update phase after the computation graph has been built.

Some readers may be confused: earlier, in "the optimization step implementation", we showed an in-place operation for the weight update, and yet the previous section said "use regular operations in the forward pass during training".

There is in fact an important point to understand properly here, so let's dig into it a little.

In-Place Operations During Training: Forward Pass vs. Weight Update

Here we will sort out a part of the training process where the use of in-place operations tends to cause confusion.

In fact, even within "training time", whether in-place operations can be used depends on the timing.

The Flow of One Training Iteration

First, let's look at the flow of a single training iteration:

# 1. Forward pass
output = model(input_data)         # Avoid in-place operations here
loss = criterion(output, target)   # Compute the loss

# 2. Backward pass
loss.backward()                    # Compute gradients

# 3. Parameter update
optimizer.step()                   # In-place operations are OK here

Across these three steps, the permissibility of in-place operations breaks down as follows:

In the forward pass, writing the following is dangerous

def forward(self, x):
    x = self.linear1(x)
    x.relu_()  # Dangerous! Avoid in-place operations at this point
    return x

On the other hand, during the weight update, an in-place operation like the following is safe

def update_weights(parameters, lr):
    with torch.no_grad():
        for param in parameters:
            if param.grad is not None:
                param.sub_(lr * param.grad)  # OK! In-place operations are fine at this point

The reason for this difference lies in the relationship between the computation graph and backpropagation

[During the forward pass (why it is dangerous)]

• At this point, the computation graph is still being built for backpropagation
• Intermediate results will be needed later for gradient computation
• If an in-place operation overwrites values, the information needed for gradient computation is lost

[During the weight update (why it is safe)]

• At this point, backpropagation has already completed
• All computation on the current graph is finished
• A new computation graph will be built in the next iteration
• Overwriting the parameters does not affect the current gradient computation

In other words, "avoid in-place operations during training" more precisely means "avoid in-place operations in the intermediate computations of the forward pass during training". Meanwhile, "in the weight update after gradient computation has completed, in-place operations can be used".

So the question of whether in-place operations are permissible should be judged not simply by "is this training or not", but from the perspective of "is the computation graph still being built, or has the computation already finished".

To deepen our understanding further, let's revisit the code from "the optimization step implementation".

def custom_sgd_update(parameters, lr):
    with torch.no_grad():
        for param in parameters:
            if param.grad is not None:
                param.sub_(lr * param.grad)  # Update weights in place

Let's recap once more why in-place operations can be used safely in this code.

First, this processing runs inside the with torch.no_grad() context. This disables autograd tracking, indicating that in the weight update phase backpropagation has already completed and there is no need to build a new computation graph.

The key fact is that at this point, all gradient computation for the current batch has finished.

There is no longer any need to reference the current parameters or intermediate results, and a new computation graph will be built for the next batch. That is why using in-place operations at this point is safe.

On the other hand, care is needed when using them in the forward pass during training. For example, the following implementation is dangerous:

# Dangerous example!
def risky_forward(self, x):
    x = self.linear1(x)
    x.relu_()  # Doing this during training is dangerous
    x = self.linear2(x)
    return x

The problem with this code is that intermediate results needed during backpropagation are lost. The following implementation is safe instead

def safe_forward(self, x):
    x = self.linear1(x)
    x = F.relu(x)  # Use the regular ReLU during training
    x = self.linear2(x)
    return x

As you can see, the timing of in-place operations is critically important.

This custom_sgd_update implementation follows the same policy as PyTorch's official optimizers: it is an example of safely using in-place operations in the weight update phase after gradient computation has completed. By using in-place operations at the right time like this, you can achieve safe training while maintaining memory efficiency.

We have gone into quite some detail here, but we hope it has deepened your understanding.

Summary: Here Is How to Put It All to Use

Based on everything so far, let's summarize practical guidelines for using in-place operations.

First, let's look at usage during inference.

For inference in production, there is no need to account for backpropagation, so you can use in-place operations aggressively.

For example, the following code works well for inference in an image recognition model

def inference(self, x):
    x = self.conv(x)
    x = self.bn(x)
    x.relu_()  # In-place operations are safe to use during inference
    x = self.pool(x)
    x.relu_()  # No problem using them multiple times
    return x

You can also make aggressive use of them in data preprocessing at inference time. For example, in image normalization you can use them as follows

def preprocess_image(image_tensor):
    image_tensor.sub_(mean).div_(std)  # Normalize in place
    image_tensor.clamp_(min=0)  # Clamp the value range directly as well
    return image_tensor

During training, on the other hand, you need to choose the timing carefully.

When updating the model's weights, in-place operations can be used safely as follows.

def update_model(self):
    with torch.no_grad():
        for param in self.parameters():
            if param.grad is not None:
                param.sub_(learning_rate * param.grad)  # OK during weight updates

In production implementations, combining this knowledge lets you build efficient systems.

For example, a production model class can be implemented like this

class ProductionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = create_backbone()
        self.is_training = False  # Flag for training/inference mode
    
    def forward(self, x):
        # During inference, implement with memory efficiency in mind
        if not self.is_training:
            x = self.preprocess_inplace(x)
            features = self.extract_features_inplace(x)
            return self.postprocess_inplace(features)
        
        # During training, use regular operations
        x = self.preprocess_standard(x)
        features = self.extract_features_standard(x)
        return self.postprocess_standard(features)

By using in-place operations appropriately, you can expect substantial gains in memory efficiency, especially for inference in production.

In deep networks such as ResNet, using in-place operations for activation functions and batch normalization can dramatically reduce the memory used by intermediate layers.

During training, meanwhile, using in-place operations only in the weight update phase after the computation graph has been built lets you achieve both safety and efficiency.

In this way, by carefully judging the use case and the timing, in-place operations become a powerful tool for memory optimization.

GPU prices tend to rise exponentially with memory capacity, and as deep learning models continue to grow in scale, the importance and usefulness of optimization techniques like these will only increase. So let's master in-place operations and put them to work!

Appendix

Commonly Used In-Place Operations

Basic arithmetic

x.add_(value) # Addition
x.sub_(value) # Subtraction
x.mul_(value) # Multiplication
x.div_(value) # Division

Neural network operations

x.relu_() # ReLU
x.sigmoid_() # Sigmoid
x.tanh_() # Tanh

Others

PyTorch provides a wide variety of in-place operators, from basic arithmetic to advanced mathematical operations.

pow_ for exponentiation and neg_ for negation are also available.

For matrix operations, there are matmul_ for matrix multiplication, transpose_ for transposing tensors, and permute_ for reordering dimensions.

For element-wise operations, there are clamp_min_ for clipping at a minimum value, clamp_max_ for clipping at a maximum value, and clamp_ for restricting the value range. minimum_ and maximum_ for element-wise minimum and maximum are also available.

For initializing or modifying data, there are zero_ to set all elements to zero, fill_ to fill with a specified value, and copy_ to copy data from another tensor. random_ for filling with random values and normal_ for filling with normally distributed random numbers are also provided.

These operators can be used just like their regular counterparts, but the trailing-underscore naming convention explicitly marks them as in-place operations. For example, while regular addition is written as tensor + value or tensor.add(value), in-place addition is written as tensor.add_(value). This makes it immediately obvious when reading code that an operation directly overwrites memory.

References

• Official PyTorch documentation
• "Deep Learning with PyTorch" ( Eli Stevens)
• "Programming PyTorch for Deep Learning"( Ian Pointer)