PyTorch
forward function
neural networks
machine learning
deep learning

What exactly does the forward function output in Pytorch?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In PyTorch, the forward method outputs whatever you explicitly return from it. In many tutorials that is a single tensor of logits or predictions, but PyTorch does not force that shape or type. A forward method can return a tensor, a tuple, a dictionary, a list of tensors, or any Python object structure that downstream code knows how to use.

What Happens When You Call a Model

When you write model(x), PyTorch does not call forward directly in the naive sense. It goes through the module call machinery, which then invokes forward and also handles hooks and other framework behavior.

A minimal model:

python
1import torch
2import torch.nn as nn
3
4class Net(nn.Module):
5    def __init__(self):
6        super().__init__()
7        self.linear = nn.Linear(4, 2)
8
9    def forward(self, x):
10        return self.linear(x)
11
12
13model = Net()
14x = torch.randn(3, 4)
15y = model(x)
16print(y.shape)

Here forward returns a tensor of shape 3 x 2, so model(x) also returns that tensor.

There Is No Special Hidden Output Type

A common misconception is that forward must return predictions in some special internal PyTorch format. It does not.

PyTorch only requires that the returned object works with the rest of your training or inference code. For example, if you plan to compute a loss, the loss function must know how to consume the returned value.

That is why many training loops assume the model returns a tensor, but that is a convention of the surrounding code, not a hard rule of forward itself.

Returning Multiple Values

It is common to return more than one value.

python
1import torch
2import torch.nn as nn
3
4class Net(nn.Module):
5    def __init__(self):
6        super().__init__()
7        self.linear = nn.Linear(4, 2)
8
9    def forward(self, x):
10        logits = self.linear(x)
11        probs = torch.softmax(logits, dim=-1)
12        return logits, probs
13
14
15model = Net()
16x = torch.randn(3, 4)
17logits, probs = model(x)
18print(logits.shape, probs.shape)

This can be useful when you want both raw logits for a loss function and normalized probabilities for reporting.

Returning Dictionaries Is Also Valid

Large models often return named outputs.

python
1import torch
2import torch.nn as nn
3
4class Net(nn.Module):
5    def __init__(self):
6        super().__init__()
7        self.linear = nn.Linear(4, 2)
8
9    def forward(self, x):
10        logits = self.linear(x)
11        return {
12            "logits": logits,
13            "predicted_class": torch.argmax(logits, dim=-1)
14        }
15
16
17model = Net()
18result = model(torch.randn(3, 4))
19print(result["logits"].shape)
20print(result["predicted_class"])

Again, this is perfectly legal as long as your training and inference code expects that structure.

The Output Is Often Logits, Not Final Class Labels

In classification examples, forward often returns logits rather than already-thresholded labels or argmax classes. That is because loss functions such as cross-entropy often expect raw logits.

So if you ask "what exactly does forward output?" the most accurate answer is:

  • whatever the model designer decided to return
  • often a tensor of logits in classification models
  • often a tensor of continuous predictions in regression models
  • sometimes a richer structure for more complex pipelines

Common Pitfalls

The biggest mistake is assuming forward must always return a single tensor.

Another mistake is returning post-processed probabilities or class indices when the loss function actually expects raw logits.

A third issue is calling model.forward(x) directly instead of model(x), which bypasses parts of the normal module call path.

Summary

  • 'forward returns whatever object you explicitly return from the method'
  • In many models that is a tensor, but it can also be a tuple, dict, or other structure
  • The surrounding training and inference code determines what output format is practical
  • Classification models often return logits, not final labels
  • Use model(x) rather than calling forward directly in normal PyTorch code

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.