Introduction

Both torch.no_grad() and torch.inference_mode() disable gradient computation in PyTorch, but they are not interchangeable. inference_mode is stricter and faster, designed for production inference paths. no_grad is more permissive, designed for situations where you want to stop gradient tracking but still need normal tensor behavior. Choosing the right one depends on whether your code is purely read-only inference or a mixed workflow.

What torch.no_grad() Does

torch.no_grad() tells PyTorch to stop building the autograd computation graph for operations inside the context. This saves memory (no gradient tensors stored) and speeds things up (no graph construction overhead).

1import torch
2
3model = MyModel()
4model.eval()
5
6x = torch.randn(1, 3, 224, 224)
7
8with torch.no_grad():
9    output = model(x)
10    # output.requires_grad is False
11    # No computation graph built
12    # But output tensors behave like normal tensors

Tensors produced inside torch.no_grad() are regular tensors. You can mutate them, use them in subsequent autograd operations outside the block, and treat them like any other tensor. The only thing no_grad removes is the gradient recording.

As a decorator

Both context managers also work as decorators:

@torch.no_grad()
def predict(model, x):
    return model(x)

What torch.inference_mode() Does

torch.inference_mode() also disables autograd, but it goes further by creating inference tensors that carry additional restrictions. These tensors cannot be used in autograd computations later, cannot be mutated in certain ways, and allow PyTorch to skip internal bookkeeping that no_grad still performs.

1import torch
2
3model = MyModel()
4model.eval()
5
6x = torch.randn(1, 3, 224, 224)
7
8with torch.inference_mode():
9    output = model(x)
10    # output is an "inference tensor"
11    # Cannot be used in autograd outside this block
12    # PyTorch skips version counting and view tracking

The key difference: tensors created inside inference_mode are marked as inference tensors. If you try to use them in a gradient-tracked operation outside the block, PyTorch raises an error.

1with torch.inference_mode():
2    y = model(x)
3
4# This will raise an error:
5z = y * torch.tensor([1.0], requires_grad=True)
6# RuntimeError: Inference tensors cannot be used in autograd

Side-by-Side Comparison

Performance Difference

The performance gap comes from the internal bookkeeping that inference_mode eliminates:

1. Version counting: PyTorch tracks a version counter on every tensor to detect in-place modifications that could invalidate the autograd graph. inference_mode skips this entirely.
2. View tracking: When you create a view of a tensor (e.g., tensor.view(...), tensor[0]), PyTorch normally tracks the relationship between the view and the base tensor for autograd correctness. inference_mode skips this.

These savings compound in models with many operations, especially transformer architectures with heavy view and reshape usage.

1import torch
2import time
3
4model = torch.nn.Linear(1000, 1000).cuda()
5x = torch.randn(64, 1000).cuda()
6
7# Benchmark no_grad
8torch.cuda.synchronize()
9start = time.perf_counter()
10for _ in range(10000):
11    with torch.no_grad():
12        _ = model(x)
13torch.cuda.synchronize()
14no_grad_time = time.perf_counter() - start
15
16# Benchmark inference_mode
17torch.cuda.synchronize()
18start = time.perf_counter()
19for _ in range(10000):
20    with torch.inference_mode():
21        _ = model(x)
22torch.cuda.synchronize()
23inference_time = time.perf_counter() - start
24
25print(f"no_grad: {no_grad_time:.3f}s")
26print(f"inference_mode: {inference_time:.3f}s")

The exact speedup varies by model and hardware. For small models the difference is negligible. For large models with complex tensor operations, inference_mode can be measurably faster.

When to Use Each One

Use torch.inference_mode() when:

• Serving predictions in a production API
• Running evaluation on a test dataset with no training interaction
• Deploying a model where the output tensors never need gradients
• Benchmarking inference throughput

1@torch.inference_mode()
2def serve_prediction(model, input_tensor):
3    model.eval()
4    return model(input_tensor)

Use torch.no_grad() when:

• Running a validation loop inside a training script where tensors might interact with the training logic
• Computing loss for logging during training (where loss tensors may be used in metrics that touch autograd)
• Debugging model outputs and inspecting intermediate tensors
• Working with older codebases that were written before inference_mode existed
• Using GAN training loops where generator outputs feed back into discriminator gradients

1# Validation loop inside training
2model.eval()
3val_loss = 0.0
4
5with torch.no_grad():
6    for batch in val_loader:
7        output = model(batch)
8        loss = criterion(output, batch.target)
9        val_loss += loss.item()
10
11model.train()

Critical: Neither Replaces model.eval()

This is the most common source of bugs. torch.no_grad() and torch.inference_mode() control autograd behavior. model.eval() controls module behavior. They solve different problems.

If your model has dropout or batch normalization layers and you only use a context manager without calling model.eval(), those layers still behave as if they are training. Dropout will randomly zero activations. BatchNorm will use batch statistics instead of running statistics. Your inference results will be wrong and non-deterministic.

Correct inference setup always combines both:

1model.eval()  # Switch modules to eval behavior
2
3with torch.inference_mode():  # Disable autograd
4    output = model(input_tensor)

Using as Function Decorators

Both work as decorators, which is cleaner for standalone prediction functions:

1@torch.inference_mode()
2def predict(model, x):
3    model.eval()
4    return model(x)
5
6@torch.no_grad()
7def compute_metrics(model, data):
8    model.eval()
9    outputs = model(data)
10    return calculate_metrics(outputs)

The decorator form is equivalent to wrapping the entire function body in the context manager.

Common Pitfalls

• Using inference_mode in a training loop's validation step when loss tensors are later used in autograd. If any tensor from inside the block touches a gradient-tracked computation, you get a RuntimeError. Use no_grad for validation loops that interact with training logic.
• Forgetting model.eval(). Neither context manager changes dropout or batch norm behavior. You must call model.eval() separately.
• Using no_grad everywhere out of habit. If your inference path is genuinely read-only, you are leaving performance on the table. Switch to inference_mode for production serving.
• Assuming the two are interchangeable. They have different guarantees. inference_mode is intentionally stricter. Code that works under no_grad may fail under inference_mode.
• Nesting contexts unnecessarily. inference_mode inside no_grad (or vice versa) works but adds no benefit. The inner context is redundant.
• Not calling model.train() after validation. If you set model.eval() for validation but forget to call model.train() before the next training iteration, your model trains with frozen dropout and batch norm.

Summary

• torch.no_grad() disables gradient tracking while keeping tensors flexible and compatible with subsequent autograd operations.
• torch.inference_mode() is a stricter, faster mode that creates inference tensors with additional restrictions, enabling PyTorch to skip version counting and view tracking.
• Use inference_mode for production inference, model serving, and pure evaluation where output tensors never re-enter autograd.
• Use no_grad for validation loops inside training, debugging, and mixed workflows where tensors might interact with gradient-tracked code.
• Always combine either context manager with model.eval() to get correct behavior from dropout and batch normalization layers.
• The performance difference is most noticeable in large models with heavy tensor operations. For small models, both are effectively equivalent.