Introduction

Predicting a single example in PyTorch uses the same model as batch prediction, but you still have to respect the model’s expected tensor shape, device, and evaluation mode. The most common mistakes are forgetting the batch dimension or leaving the model in training mode. Once those are handled, single-example inference is straightforward.

Put the Model in Evaluation Mode

Before inference, switch the model to evaluation mode.

1import torch
2import torch.nn as nn
3
4model = nn.Sequential(
5    nn.Linear(4, 8),
6    nn.ReLU(),
7    nn.Linear(8, 3)
8)
9
10model.eval()

model.eval() matters because layers such as dropout and batch normalization behave differently during training and inference.

If you loaded trained weights, do that first:

model.load_state_dict(torch.load("model.pth", map_location="cpu"))
model.eval()

Add a Batch Dimension

PyTorch models usually expect input with a batch dimension, even if you only want one sample.

1sample = torch.tensor([5.1, 3.5, 1.4, 0.2], dtype=torch.float32)
2sample = sample.unsqueeze(0)
3
4print(sample.shape)  # torch.Size([1, 4])

Without unsqueeze(0), the tensor shape is just [4], which often does not match what the network expects.

Run Inference with torch.no_grad()

Wrap prediction in torch.no_grad() so PyTorch does not build a gradient graph unnecessarily.

1with torch.no_grad():
2    logits = model(sample)
3    predicted_class = torch.argmax(logits, dim=1)
4
5print(predicted_class.item())

This is the standard pattern for inference:

1. prepare the tensor
2. switch the model to eval mode
3. call the model inside torch.no_grad()
4. post-process the output

Convert Raw Output into Probabilities

For classification, model outputs are often logits rather than human-readable probabilities.

1with torch.no_grad():
2    logits = model(sample)
3    probabilities = torch.softmax(logits, dim=1)
4    predicted_class = torch.argmax(probabilities, dim=1)
5
6print(probabilities)
7print(predicted_class.item())

Use softmax only if you actually want probabilities for reporting or thresholding. If you only need the top class, argmax on logits is enough.

Example with an Image Tensor

For image models, the same rules apply but the shape is usually [batch, channels, height, width].

1import torch
2
3image = torch.randn(3, 224, 224)
4image = image.unsqueeze(0)
5
6model.eval()
7with torch.no_grad():
8    output = model(image)
9    pred = output.argmax(dim=1)
10
11print(pred.item())

The extra batch dimension is still required even for one image.

In real code, do not forget the same preprocessing used during training, such as resize, normalization, and channel order.

Match Device and Dtype

If the model is on GPU, the input must also be on GPU.

1device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
2model.to(device)
3sample = sample.to(device)
4
5with torch.no_grad():
6    logits = model(sample)

The same applies to data types. A model trained on float32 inputs should usually receive float32 tensors at inference time.

If you loaded label names separately, this is also the stage where you map the predicted class index back to a human-readable label before returning or displaying the result.

Common Pitfalls

The most common mistake is forgetting the batch dimension. A single sample still usually needs shape [1, features] or [1, channels, height, width].

Another issue is skipping model.eval(), which can make predictions unstable when dropout or batch normalization is involved.

Some developers also run inference without torch.no_grad(). The code still works, but it wastes memory and computation by tracking gradients that are never used.

Finally, make sure preprocessing at prediction time matches training. If the model was trained on normalized inputs and you feed raw values, the prediction can be wrong even though the code runs.

Summary

• Call model.eval() before single-example inference.
• Add a batch dimension with unsqueeze(0).
• Use torch.no_grad() during prediction.
• Apply argmax or softmax depending on how you want to interpret the output.
• Keep tensor shape, device, dtype, and preprocessing consistent with training.