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.
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:
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.
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.
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=0is 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
- How to continue training model using ModelCheckpoint of Keras
- How to control GPU memory size with tf.estimator
- How to control memory while using Keras with tensorflow backend?
- How to convert a list of tensors of dim N to a tensor of dim N1
- How to convert a Hugging Face Pytorch model AutoTrain to TorchScript .pt for deployment?
- How to convert a PyTorch nn.Module into a HuggingFace PreTrainedModel object?
- How to continue to train SVM based on the previous model
- how to control frequency of loss logging messages when using tf.Estimator?
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.