PyTorch
neural network
machine learning
programming tutorial
two inputs

How to construct a network with two inputs 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

Building a PyTorch model with two inputs is a normal pattern when different pieces of information should be processed in different ways before being combined. A common example is an image plus metadata, or a text embedding plus numeric features from a database row.

Use Separate Branches Before Merging

The cleanest design is usually a branch for each input type, followed by a merge step such as concatenation. Each branch learns a representation suited to its own input, and the merged tensor is then passed to a shared classifier or regressor.

Here is a runnable example with an image-like tensor and a small metadata vector:

python
1import torch
2import torch.nn as nn
3
4class TwoInputNet(nn.Module):
5    def __init__(self):
6        super().__init__()
7
8        self.image_branch = nn.Sequential(
9            nn.Linear(28 * 28, 128),
10            nn.ReLU(),
11            nn.Linear(128, 64),
12            nn.ReLU(),
13        )
14
15        self.meta_branch = nn.Sequential(
16            nn.Linear(5, 16),
17            nn.ReLU(),
18        )
19
20        self.classifier = nn.Sequential(
21            nn.Linear(64 + 16, 32),
22            nn.ReLU(),
23            nn.Linear(32, 10),
24        )
25
26    def forward(self, image, meta):
27        image = image.view(image.size(0), -1)
28        image_features = self.image_branch(image)
29        meta_features = self.meta_branch(meta)
30
31        combined = torch.cat([image_features, meta_features], dim=1)
32        return self.classifier(combined)
33
34model = TwoInputNet()
35
36image_batch = torch.randn(8, 1, 28, 28)
37meta_batch = torch.randn(8, 5)
38
39logits = model(image_batch, meta_batch)
40print(logits.shape)

The important detail is dim=1 in torch.cat. That concatenates features across the channel or feature dimension while preserving the batch size.

Match the Data Loader to the Forward Signature

Once the model accepts two tensors, the training loop must provide two tensors in the same order. A custom dataset can return a tuple like (image, meta, target), and the training loop passes the first two items into the model.

python
1import torch
2from torch.utils.data import Dataset, DataLoader
3
4class DemoDataset(Dataset):
5    def __len__(self):
6        return 100
7
8    def __getitem__(self, index):
9        image = torch.randn(1, 28, 28)
10        meta = torch.randn(5)
11        target = torch.randint(0, 10, size=(1,)).item()
12        return image, meta, target
13
14dataset = DemoDataset()
15loader = DataLoader(dataset, batch_size=16, shuffle=True)
16
17for images, meta, targets in loader:
18    logits = model(images, meta)
19    print(logits.shape, targets.shape)
20    break

This pattern scales well because the dataset stays responsible for packaging related inputs together.

Training Looks Almost the Same

The optimization step is not fundamentally different from a single-input model. You only need to unpack both inputs and send them to the same device.

python
1device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
2model = model.to(device)
3
4criterion = nn.CrossEntropyLoss()
5optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
6
7for images, meta, targets in loader:
8    images = images.to(device)
9    meta = meta.to(device)
10    targets = targets.to(device)
11
12    optimizer.zero_grad()
13    logits = model(images, meta)
14    loss = criterion(logits, targets)
15    loss.backward()
16    optimizer.step()
17    break

The key idea is that PyTorch does not need a special multi-input API. A model can accept as many tensors as its forward method defines.

When Inputs Need Different Architectures

In the example above, both branches use fully connected layers. In practice, you often mix architectures:

  • use a CNN branch for images
  • use an embedding or recurrent branch for text
  • use a small multilayer perceptron for numeric metadata

The merge step stays the same. Each branch transforms its input into a feature vector, and the shared head learns how to combine those vectors for the final task.

If the two inputs have very different scales or importance, normalize them appropriately before training. The branch design matters more than the fact that there are two inputs.

Common Pitfalls

  • Concatenating on the wrong dimension. If dim=0 is used by mistake, batch items get mixed together.
  • Forgetting to flatten the image branch before sending it into linear layers.
  • Returning mismatched batch sizes from the dataset, which makes concatenation fail immediately.
  • Moving one input tensor to the GPU and leaving the other on the CPU.
  • Designing one branch to output a huge feature vector while the other branch is tiny, which can drown out the smaller signal.

Summary

  • Multi-input PyTorch models are built by defining multiple arguments in forward.
  • Separate branches usually make the model easier to reason about and train.
  • Merge the branch outputs with torch.cat(..., dim=1) after they are shaped as feature vectors.
  • The training loop stays simple: unpack both inputs, move both to the device, and compute loss normally.
  • Most bugs come from shape mismatches, incorrect concatenation dimensions, or inconsistent dataset output.

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.