Want to ace your PyTorch interview?
Well, good news! In this guide, I’m breaking down some of the most important PyTorch interview questions you might get asked, covering everything from tensors and autograd through to training models, GPU performance, and more advanced topics.
Better still, rather than just giving you definitions to memorize, I’ll break down what each concept actually means and why it matters, so you can confidently explain the answer in your own words.
Get through these and you'll go into your interview feeling prepared and ready to land that job.
Let's dive in...
Sidenote: If you find that you struggle to answer any of these questions and want to brush up, or simply work on some new projects to flesh out your portfolio, then check out my PyTorch course:
This PyTorch course is your step-by-step guide to developing your own deep learning models using PyTorch. You'll learn Deep Learning with PyTorch by building a massive 3-part real-world milestone project. By the end, you'll have the skills and portfolio to get hired as a Deep Learning Engineer.
With that out of the way, let’s get into the questions...
Beginner Pytorch interview questions
#1. What is PyTorch, and why would you use it?
PyTorch is an open-source machine learning framework that's commonly used to build and train neural networks.
The easiest way to understand why you'd use it is to think about everything involved in training a machine learning model.
For example
Suppose you wanted to build a neural network that recognizes handwritten numbers.
At a very simplified level, you need to:
Load your training data
↓
Pass that data through the model
↓
Measure how wrong its predictions are
↓
Calculate how each parameter contributed to that error
↓
Update those parameters
↓
RepeatYou could implement all of the mathematics and training logic yourself, but that would quickly become complicated.
PyTorch provides the building blocks that handle much of this for you.
For example
Tensors let you efficiently work with the multidimensional numerical data used by your model. PyTorch's autograd system can automatically calculate gradients, which you need when training neural networks. While the torch.nn module gives you reusable components for building the network itself, while optimizers help update its parameters during training.
PyTorch can also perform computations on hardware accelerators such as GPUs, which becomes extremely important when you're training larger models or working with large datasets.
One reason developers like PyTorch is that it also feels quite natural to write in Python.
For example
A small neural network can look something like this:
import torch.nn as nn
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.network = nn.Sequential(
nn.Linear(784, 128),
nn.ReLU(),
nn.Linear(128, 10)
)
def forward(self, x):
return self.network(x)
The important thing for now is that PyTorch gives you a framework for defining the model, performing the computations needed to train it, calculating gradients, updating its parameters, and running those computations efficiently on suitable hardware.
So in an interview, you could explain it like this:
"PyTorch is an open-source machine learning framework commonly used for building and training neural networks. It provides things like tensors, automatic differentiation, neural network components, optimizers, data utilities, and hardware acceleration, so you don't have to implement the underlying training infrastructure yourself.
I'd use it when I need the flexibility to build, train, and experiment with machine learning models in Python."
The important thing to understand is that PyTorch isn't the neural network itself. It's the framework that gives you the tools to build and train one.
And the most fundamental of those tools is something we're going to use throughout this guide, which is tensors.
#2. What is a tensor in PyTorch?
Tensors are the fundamental data structure used in PyTorch.
If you've worked with Python before, you've probably used lists. And if you've done any data science or numerical computing, you may have used NumPy arrays.
A tensor is similar to a NumPy array in that it stores numerical data across one or more dimensions.
For example
You could have a single number: 5
A one-dimensional tensor: [1, 2, 3]
A two-dimensional tensor:
[
[1, 2, 3],
[4, 5, 6]
]Or tensors with three, four, or many more dimensions.
In PyTorch, you can create them using torch.tensor():
import torch
x = torch.tensor([
[1, 2, 3],
[4, 5, 6]
])
print(x.shape)Or in other words, it has two rows and three columns.
But why does machine learning need tensors? Simply because almost everything your model works with eventually needs to be represented numerically.
An image, for example, might be represented as a three-dimensional tensor: channels × height × width.
While a batch of those images adds another dimension: batch × channels × height × width.
So 32 RGB images that are each 224 × 224 pixels might have a shape like: [32, 3, 224, 224].
Text works this way too.
For example
When you give a large language model (LLM) text or code, a tokenizer breaks that input into smaller pieces called tokens. These might represent whole words, parts of words, or punctuation. Each token has a numerical ID, and those IDs can be stored in a tensor for the model to process.
During generation, the model predicts which tokens could come next, building its response one token at a time. The generated token IDs are then decoded back into readable text or code.
The same idea applies to other types of data. Model inputs, outputs, weights, gradients, and intermediate calculations can all be represented using tensors.
Where PyTorch tensors become particularly useful is that they're designed for machine learning workloads.
They can perform highly optimized mathematical operations, run computations on accelerators such as GPUs, and integrate with PyTorch's automatic differentiation system so gradients can be calculated during training.
For example
You can create a tensor on a GPU and perform operations there rather than on the CPU with x = torch.tensor([1.0, 2.0, 3.0], device="cuda").
Assuming a CUDA-compatible GPU is available, the tensor and operations involving it can now live on the GPU.
So in an interview, you could explain it like this:
"A tensor is PyTorch's fundamental data structure for representing numerical data. It's similar to a multidimensional NumPy array, but PyTorch tensors are designed to integrate with features such as GPU acceleration and automatic differentiation. They're used throughout a PyTorch application for things like model inputs, parameters, outputs, and gradients."
The easiest mental model is you find that hard to remember, is that a tensor is the numerical container PyTorch uses to perform its computations.
And because they look quite similar to NumPy arrays, that raises an obvious interview question: what's actually different about them?
#3. How are PyTorch tensors different from NumPy arrays?
At first glance, PyTorch tensors and NumPy arrays look very similar.
For example
Here's a NumPy array:
import numpy as np
x = np.array([1, 2, 3])And here's roughly the same thing as a PyTorch tensor:
import torch
x = torch.tensor([1, 2, 3])
So why does PyTorch need tensors at all?
Well, there are two particularly important differences for machine learning.
The first is hardware acceleration.
A normal NumPy array generally performs its computations on the CPU. PyTorch tensors can be moved to supported accelerators such as a GPU, allowing PyTorch to perform the large numbers of mathematical operations involved in training neural networks much more efficiently.
For example
x = torch.tensor([1.0, 2.0, 3.0])
if torch.cuda.is_available():
x = x.to("cuda")Now x can participate in computations on the CUDA GPU rather than the CPU.
The second major difference is automatic differentiation.
When training a neural network, you need to calculate how changes to the model's parameters would affect its error. Those derivatives are called gradients, and PyTorch can automatically calculate them using its autograd system.
For example
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
y.backward()
print(x.gradThe result is tensor(4.).
Why 4?
Because:
y = x²
dy/dx = 2x
x = 2
2 × 2 = 4By setting requires_grad=True, PyTorch tracks the operations involving x. When we call backward(), it can work backwards through those operations and calculate the gradient.
NumPy isn't designed to provide that automatic differentiation system.
That's a huge distinction because neural networks can contain millions or even billions of parameters. You obviously don't want to manually derive and implement the gradient calculation for every operation in the model.
PyTorch tensors and NumPy arrays can also work together.
For example
You can convert a NumPy array into a tensor:
array = np.array([1, 2, 3])
tensor = torch.from_numpy(array)And convert a CPU tensor to a NumPy array with array = tensor.numpy().
One detail worth knowing is that these conversions can sometimes share the same underlying memory rather than copying the data. That means modifying one can also change the other.
So in an interview, you could explain it like this:
"PyTorch tensors are similar to NumPy arrays because both represent multidimensional numerical data and support many of the same mathematical operations. The big differences for machine learning are that PyTorch tensors can run computations on hardware accelerators such as GPUs and integrate with autograd for automatic differentiation. PyTorch also makes it easy to convert between tensors and NumPy arrays."
The important distinction isn't that tensors are a completely different kind of data structure. It's that PyTorch tensors add the capabilities needed to build and train neural networks.
And we've just introduced one of the most important of those capabilities which is autograd.
#4. What is autograd, and how does automatic differentiation work in PyTorch?
When you're training a neural network, you need to figure out how each of the model's parameters contributed to the error it made. That's where gradients come in.
For example
Imagine an extremely simple model with a single weight:
Input
↓
Weight
↓
Prediction
↓
LossThe loss tells us how wrong the prediction was, but knowing that the model was wrong isn't enough. We need to know how we should change the weight to make the loss smaller.
Mathematically, that means calculating the derivative of the loss with respect to the weight:
dLoss
─────
dWeightFor a real neural network, however, the output might depend on millions of parameters connected through many layers and mathematical operations.
Calculating all of those derivatives manually would be incredibly difficult but PyTorch's autograd system does it automatically.
For example
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
z = y * 3
z.backward()
print(x.grad)This gives us tensor(12.).
So where did 12 come from? Well, our calculation was:
y = x²
z = 3yTherefore z = 3x², because the derivative of z with respect to x is dz/dx = 6x. And because x is 2, that means that 6 × 2 = 12
Yes this is super simple math but PyTorch calculated that for us. The most important part to look at is how it did that, because we created x with requires_grad=True.
This works because PyTorch tracks the operations involving that tensor as they're performed.
For example
You can think of it as building a computational graph:
x = 2
↓
x²
↓
y
↓
× 3
↓
zSo, when we call z.backward(), autograd works backwards through that graph, applying the chain rule to calculate how z depends on the tensors that require gradients.Then the resulting gradient for x is then stored in X.grad.
Make sense?
Well, this same basic process is what happens during neural network training.
The graph is obviously much larger, but conceptually it's still:
Model parameters
↓
Forward pass
↓
Predictions
↓
Loss
↓
loss.backward()
↓
Gradients for parametersAn optimizer can then use those gradients to update the model's parameters in a direction intended to reduce the loss.
There's another detail that's worth understanding though, and that's the fact that PyTorch builds this graph dynamically as operations execute.
This means the graph reflects the computations that actually happened during that particular forward pass. This is one reason PyTorch feels natural to work with in Python, because you can use ordinary Python control flow while defining model behavior.
So in an interview, you could explain it like this:
"Autograd is PyTorch's automatic differentiation system. When operations involve tensors that require gradients, PyTorch records those operations in a computational graph. When you call backward() on the result, autograd works backwards through that graph using the chain rule and calculates gradients for the relevant tensors. Those gradients are then used by optimizers during training to update the model's parameters."
The main idea is that the forward pass calculates what the model predicts. Autograd helps calculate how the parameters contributed to the resulting loss so we know how to update them.
We've used requires_grad=True a couple of times now, which makes the next question a natural one…
#5. What does requires_grad do in PyTorch?
We’ve already used requires_grad=True in the last couple of questions, so let’s look at what it actually does.
By default, PyTorch doesn’t need to track every operation you perform on every tensor.
For example
import torch
x = torch.tensor(2.0)
y = x ** 2You can see this by checking print(x.requires_grad), which returns False.
However, suppose x represents a value that we want to optimize. Well then we can tell PyTorch to track operations involving it by setting x = torch.tensor(2.0, requires_grad=True).
So now y = x ** 2.
This then becomes part of the computation autograd is tracking, so if we call y.backward(), PyTorch can calculate the gradient and store it in X.grad, which gives us tensor(4.).
This becomes especially important when you're building neural networks.
Why?
Well, a model might contain millions (or billions!) of parameters, such as weights and biases, that need to be adjusted during training.
PyTorch's nn.Module handles much of this for you. Parameters created by modules such as nn.Linear normally have requires_grad=True, which means autograd tracks the operations needed to calculate their gradients.
You can see that with a simple layer:
import torch.nn as nn
layer = nn.Linear(3, 1)
print(layer.weight.requires_grad)However, sometimes you don't want particular parameters to change.
For example
Imagine you're fine-tuning a pretrained model, but you only want to train its final layer. You could freeze the parameters in the earlier layers like so:
for param in model.parameters():
param.requires_grad = Falsefor param in model.classifier.parameters():
param.requires_grad = True(One subtle point is that requires_grad=False doesn't mean you can't use that tensor in a calculation. It simply means you aren't asking PyTorch to accumulate a gradient for that tensor as something you want to optimize).
So in an interview, you could explain it like this:
"requires_grad tells PyTorch whether a tensor should participate in gradient tracking for automatic differentiation. Model parameters that need to be learned normally have it set to True, which allows autograd to calculate their gradients during backpropagation. You can set it to False for parameters you don't want to train, such as when freezing part of a pretrained model."
The easiest way to remember it is:
requires_grad=True
↓
Track the operations needed for gradients
↓
backward()
↓
Gradient available for the parameter
↓
Optimizer can use it during trainingWe've now covered tensors and how PyTorch calculates gradients with them. The next step is to see how PyTorch organizes all those trainable parameters into an actual model.
That's where
nn.Module comes in.
#6. What is nn.Module, and why do PyTorch models inherit from it?
So far, we've looked at individual tensors and gradients, but a real neural network might contain dozens or even hundreds of layers, each with its own weights and biases that need to be trained.
PyTorch needs a convenient way to organize all of that, and that's what the nn.Module is for.
nn.Module is the base class used to build neural networks and other reusable model components in PyTorch.
For example
Here's the simple network we saw earlier:
import torch
import torch.nn as nn
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(784, 128)
self.relu = nn.ReLU()
self.layer2 = nn.Linear(128, 10)
def forward(self, x):
x = self.layer1(x)
x = self.relu(x)
x = self.layer2(x)
return xThere are two important parts to understand here.
First, the __init__() method is where we define the layers and other components our model needs:
self.layer1 = nn.Linear(784, 128)
self.relu = nn.ReLU()
self.layer2 = nn.Linear(128, 10)Then forward() defines how data should move through those components:
Input
↓
Linear layer
↓
ReLU
↓
Linear layer
↓
OutputSo if we create the model model = NeuralNetwork(), and give it some input:
x = torch.randn(1, 784)
output = model(x)Make sense?
However, you might notice that we wrote model(x) rather than model.forward(x).
This was intentional.
Why?
Well, when an nn.Module is called, PyTorch handles the call and then invokes forward() internally. But, you should normally call the module itself rather than calling forward() directly because PyTorch's module call machinery also handles features such as hooks.
The good news is that nn.Module does more than give us somewhere to put a forward() method.
One of its most useful jobs is keeping track of the model's parameters and submodules.
For example
nn.Linear contains trainable weights and biases. Because we assigned those layers as attributes of our module, PyTorch automatically registers them as part of the model.
That means we can do things like:
for parameter in model.parameters():
print(parameter.shape)optimizer = torch.optim.Adam(
model.parameters(),
lr=0.001
)The optimizer now knows which parameters it needs to update during training, so the module can also keep track of its state, move its parameters between devices, switch between training and evaluation behavior, and be composed from other modules.
That last point is particularly useful, because an entire neural network can be an nn.Module, but so can a single layer or a reusable block inside a much larger model.
You can essentially build models out of smaller modules like so:
Model
├── Feature extractor
│ ├── Layer
│ ├── Layer
│ └── Layer
│
└── Classifier
├── Layer
└── LayerEach component can be an nn.Module, and PyTorch keeps track of how they're nested together.
So in an interview, you could explain it like this:
"nn.Module is the base class for neural network models and components in PyTorch. You define your layers or submodules in __init__() and define the forward computation in forward(). Inheriting from nn.Module also lets PyTorch automatically register parameters and submodules, which makes it possible to train, save, move, and manage the model."
The important idea is that nn.Module gives PyTorch a structured way to understand your model and everything inside it.
But having a model isn't enough. We also need an efficient way to get training data into it, which brings us to Dataset and DataLoader.
#7. What are Dataset and DataLoader, and how do they work together?
Once you've built a model, you need a way to get your training data into it.
For a tiny dataset, you could potentially load everything into memory yourself and pass tensors directly to the model. But real datasets can contain thousands or millions of examples.
You normally don't want to load the entire dataset into your model at once.So instead, training usually happens in smaller batches.
For example
10,000 training examples
↓
Split into batches of 32
↓
Batch 1 → Model
Batch 2 → Model
Batch 3 → Model
...PyTorch provides two important abstractions for handling this: Dataset and DataLoader.
A Dataset represents the data itself. Its job is essentially to answer two questions:
How many examples are there?
Give me example #n
For example
For a custom dataset, you can define those behaviors using __len__() and __getitem__():
from torch.utils.data import Dataset
class MyDataset(Dataset):
def __init__(self, features, labels):
self.features = features
self.labels = labels
def __len__(self):
return len(self.features)
def __getitem__(self, index):
return self.features[index], self.labels[index]Now we can retrieve an individual example with features, label = dataset[0].
Handy right?
However, even though the Dataset knows how to access one sample, during training we usually want to do more than retrieve examples individually. That's where DataLoader comes in.
A DataLoader takes a dataset and handles the process of feeding that data to your training loop.
For example
from torch.utils.data import DataLoader
train_loader = DataLoader(
dataset,
batch_size=32,
shuffle=True
)for features, labels in train_loader:
# Train the model using this batch
PassThe relationship looks something like this:
Dataset
↓
Individual samples
↓
DataLoader
↓
Batch samples together
↓
Optionally shuffle them
↓
Training loopThe batch_size=32 argument tells the loader to combine samples into batches of 32.
And shuffle=True means the training examples are reshuffled between epochs rather than always being presented in the same order.
DataLoader can also do more than batching and shuffling. For example, it can use worker processes to load and prepare data in parallel, which can become important when loading the data itself is slowing down training.
So the distinction between the two is fairly simple. Dataset defines how your individual examples are stored or accessed, while DataLoader handles how those examples are delivered to your training loop.
So in an interview, you could explain it like this:
"Dataset represents a collection of training examples and defines how individual samples are accessed, usually through __len__() and __getitem__(). A DataLoader wraps that dataset and handles things like batching, shuffling, and parallel data loading so the data can be efficiently fed into the training loop."
An easy way to remember the relationship is:
Dataset = What data do I have?
DataLoader = How do I feed that data to the model?
Intermediate PyTorch interview questions
#8. Walk me through a typical PyTorch training loop
A training loop is the process PyTorch repeats to gradually improve a model's parameters using the training data.
At a high level, each batch goes through the same basic steps:
Get a batch of data
↓
Run the model
↓
Calculate the loss
↓
Calculate gradients
↓
Update the parameters
↓
RepeatFor example
Suppose we've already created our model, chosen a loss function, and created an optimizer:
import torch.nn as nn
import torch.optim as optim
model = NeuralNetwork()
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(
model.parameters(),
lr=0.001
)The loss function measures how far the model's predictions are from the correct answers
The optimizer is responsible for updating the model's parameters based on the gradients we calculate during training
Now we can create our training loop:
for epoch in range(5):
model.train()
for features, labels in train_loader:
optimizer.zero_grad()
predictions = model(features)
loss = loss_fn(predictions, labels)
loss.backward()
optimizer.step()There's quite a lot happening in those few lines, so let's walk through one batch. First optimizer.zero_grad() clears gradients left over from the previous training step.
I'll explain why that's necessary in the next question, but for now just remember that PyTorch accumulates gradients by default, so we normally clear them before calculating the next batch's gradients.
Then we perform the forward pass predictions = model(features). The input data travels through the model and produces predictions.
Next, we compare those predictions with the correct answers using loss = loss_fn(predictions, labels). This gives us a loss value representing how wrong the model currently is. Then comes the autograd system we learned about earlier loss.backward().
PyTorch works backwards through the computational graph and calculates gradients for the model's trainable parameters.
At this point, we know how those parameters contributed to the loss, but we haven't actually changed them yet. That's the optimizer's job optimizer.step().
The optimizer uses the gradients to update the model's parameters according to its optimization algorithm.
So one training step looks like this:
Input batch
↓
Forward pass
↓
Predictions
↓
Calculate loss
↓
Backward pass
↓
Gradients
↓
Optimizer step
↓
Updated parametersSo, if we train for five epochs for epoch in range(5), the model gets five passes through the training dataset, with its parameters gradually being updated along the way.
However, in a real training loop, you'd normally do more than this.
For example
You might move the data to a GPU, track the average loss, evaluate the model on validation data, save checkpoints, adjust the learning rate, or stop training when performance stops improving, but the fundamental training process remains the same.
So in an interview, you could explain it like this:
"In a typical PyTorch training loop, I iterate over batches from a DataLoader. For each batch, I clear the previous gradients, run a forward pass through the model, calculate the loss, call loss.backward() to calculate gradients using autograd, and then call optimizer.step() to update the model's parameters. I repeat that process across the dataset for however many epochs I'm training."
If you understand that sequence, you understand the basic mechanics behind training a PyTorch model.
However, there are three lines in that loop that interviewers often ask about specifically:
optimizer.zero_grad()
loss.backward()
optimizer.step()#9. What do optimizer.zero_grad(), loss.backward(), and optimizer.step() do?
optimizer.zero_grad(), loss.backward(), and optimizer.step() each perform a different part of the learning process and are at the heart of most PyTorch training loops.
Let's start with loss.backward(). We touched on a little when we looked at autograd.
So, after the forward pass and loss calculation, PyTorch has a computational graph showing how the loss depends on the model's trainable parameters. And calling loss.backward() works backwards through that graph and calculates the gradient of the loss with respect to each parameter.
So what does that mean?
Well basically we're asking for each parameter in the model, "If I changed this slightly, how would it affect the loss?"
Those gradients are then stored in the .grad attribute of each parameter.
For example
print(model.layer1.weight.grad) would give us the gradients calculated for that layer's weights. But calculating the gradients doesn't actually change the weights.
That's what optimizer.step() does.
The optimizer looks at the gradients and uses them to update the parameters according to whatever optimization algorithm we've chosen.
With simple gradient descent, you can think of the update roughly like this new weight = old weight - learning rate × gradient.
If we're using something like Adam, the actual update is more sophisticated, but the general idea is the same in that the optimizer uses information from the gradients to decide how the parameters should change.
That leaves optimizer.zero_grad().
Why do we need to clear the gradients every time?
Because PyTorch accumulates gradients by default.
For example
Suppose the first batch produces a gradient of 2 and the second produces 3. If we simply call backward() again without clearing the previous gradient, we'd accumulate them like so 2 + 3 = 5.
Now sure, sometimes we deliberately want to accumulate gradients.
This technique is called gradient accumulation, and it lets us combine gradients from several smaller batches before updating the model's parameters. That gives us a larger effective batch size without needing enough GPU memory to process the larger batch all at once.
For example
We could process four batches of 8 examples each, scaling the loss appropriately and accumulating gradients before calling optimizer.step(). That gives us an effective batch size of 32.
In that setup, we'd clear the gradients before each new group of four batches.
But in our training loop here, we're updating the parameters after every batch. So we use optimizer.zero_grad() to clear the old gradients before calculating the next batch's gradients.
Then the sequence makes sense:
optimizer.zero_grad()
↓
Clear old gradientsForward pass + loss
↓
Calculate how wrong we areloss.backward()
↓
Calculate new gradientsoptimizer.step()
↓
Update model parametersYou may also see code that uses optimizer.zero_grad(set_to_none=True) or model.zero_grad(). But the underlying reason is the same in that you need to manage the gradients from previous backward passes rather than accidentally carrying them into the next update.
It's also worth noticing that the order matters.
So if you did this:
loss.backward()
optimizer.zero_grad()
optimizer.step()You'd calculate the gradients and then clear them before the optimizer could use them, which is not particularly helpful!
So in an interview, you could explain it like this:
"loss.backward() uses autograd to calculate gradients of the loss with respect to the model's trainable parameters. optimizer.step() uses those gradients to update the parameters, and optimizer.zero_grad() clears previously accumulated gradients so they don't unintentionally carry over into the next training step."
An easy way to remember the three is:
zero_grad() → Clear
backward() → Calculate
step() → UpdateOnce you've trained the model, though, you don't always want it behaving exactly as it did during training, which brings us to our next question.
#10. What's the difference between model.train() and model.eval()?
When you're training a model and when you're using it to make predictions, you usually want slightly different behavior.
That's why PyTorch modules have two modes: model.train() and model.eval().
However, the names can be slightly misleading at first because calling model.train() doesn't actually train the model, and calling model.eval() doesn't evaluate its accuracy. Instead, they tell certain layers in the model whether they should behave in training mode or evaluation mode.
Why does that matter?
Because some layers behave differently depending on whether you're training the model or using it for inference.
Two common examples are Dropout and Batch Normalization.
Let's start with Dropout.
During training, Dropout randomly sets some activations to zero. The idea is to stop the network from relying too heavily on particular neurons, which can help reduce overfitting.
So with model.train(), a Dropout layer is active.
But imagine leaving that randomness turned on when you're making predictions for real users. Well, the same input could potentially travel through a slightly different network each time. So when you switch to model.eval(), Dropout stops randomly dropping activations.
Batch Normalization also changes its behavior.
During training, BatchNorm updates and uses statistics based on the training process. During evaluation, it uses the running statistics it learned during training instead of continuing to update them from the current batch.
So you can think of the distinction roughly like this:
model.train()
Dropout → active
BatchNorm → training behaviorversus:
model.eval()
Dropout → disabled
BatchNorm → evaluation behaviorFor example
A typical training and validation workflow might therefore look like:
for epoch in range(5):
model.train()
for features, labels in train_loader:
optimizer.zero_grad()
predictions = model(features)
loss = loss_fn(predictions, labels)
loss.backward()
optimizer.step()
model.eval()
# Run validation…There's an important catch here, though because model.eval() does not turn off gradient calculation.
That's a really common misconception because it changes the behavior of modules that have different training and evaluation modes, but autograd can still track computations. So during evaluation or inference, you'll often see something like:
model.eval()
with torch.no_grad():
predictions = model(features)model.eval() changes the model's training-dependent behavior, while torch.no_grad() tells PyTorch not to record operations for gradient calculation. Similarly, calling model.train() doesn't somehow switch autograd on as gradient tracking is controlled separately.
So in an interview, you could explain it like this:
"model.train() and model.eval() switch an nn.Module between training and evaluation modes. This matters for layers whose behavior changes between the two, such as Dropout and BatchNorm. model.eval() doesn't disable gradient calculation, though, so for inference I'd normally use it together with something like torch.no_grad() or inference mode."
The key distinction is:
train() / eval()
↓
Control model behaviorGradient modes
↓
Control autograd behaviorKeeping those two ideas separate is important.
#11. What's the difference between torch.no_grad() and torch.inference_mode()?
We just saw that model.eval() changes how certain layers behave during evaluation, but it doesn't stop PyTorch from tracking gradients.
That's where torch.no_grad() and torch.inference_mode() come in.
Normally, when PyTorch performs operations involving tensors that require gradients, autograd keeps track of the information it may need for a later backward pass.
That's essential during training:
Forward pass
↓
Track operations
↓
Calculate loss
↓
backward()
↓
Calculate gradientsBut suppose you've finished training and you're using the model to make predictions. Well, you're not going to call backward(), so keeping all that information around is unnecessary, so what can we do?
One option is torch.no_grad():
model.eval()
with torch.no_grad():
predictions = model(features)Inside that context, PyTorch doesn't record operations for the backward graph. That reduces the overhead associated with autograd and can reduce memory usage during inference.
You might also see:
model.eval()
with torch.inference_mode():
predictions = model(features)The difference is that inference_mode() goes further.
Because PyTorch knows the tensors created inside that context are being used for inference, it can disable some additional autograd-related tracking and bookkeeping.
That can make inference more efficient than no_grad(). However, there's a trade-off, which is why torch.no_grad() is the more flexible option.
Why?
Well, tensors and computations from a no_grad() context can later be used in gradient-tracked computations.
However, inference mode is more restrictive because tensors created in inference mode can't simply be taken out and used as part of computations that autograd needs to record for a later backward pass.
So a useful mental model is:
torch.no_grad()
↓
Don't record gradients
↓
More flexible
versus:
torch.inference_mode()
↓
Don't record gradients
+
Disable additional autograd overhead
↓
Potentially faster, but more restrictiveFor pure model inference, inference_mode() is therefore a good option when your code is compatible with its restrictions.
If you need to temporarily disable gradient tracking but then use the resulting values in subsequent gradient-tracked calculations, no_grad() may be more appropriate.
And remember that neither replaces model.eval(). Which is why you may still need both because they're solving different problems.
model.eval()
with torch.inference_mode():
predictions = model(features)model.eval() changes the behavior of modules such as Dropout and BatchNorm, while torch.inference_mode() changes how autograd handles the computation.
So in an interview, you could explain it like this:
"torch.no_grad() and torch.inference_mode() both disable gradient recording, which reduces unnecessary autograd overhead when gradients aren't needed. Inference mode goes further by disabling additional autograd bookkeeping, so it can provide better performance for pure inference, but it's also more restrictive.
I'd typically use inference mode for straightforward model inference and no_grad when I need more flexibility."
The important distinction is that there are really two separate questions when you're running a trained model:
How should the model's layers behave?
↓
model.eval()
Do I need autograd for these computations?
↓
no_grad() or inference_mode()#12. How do you move a PyTorch model and tensors between CPU and GPU?
One of the big advantages of PyTorch tensors is that computations aren't limited to the CPU.
Why care?
Simply because training neural networks involves a huge number of mathematical operations, many of which can be performed much faster on a GPU.
However, there's an important rule in PyTorch, and that's that the model and the tensors involved in an operation generally need to be on the same device.
For example
Let's say we create a model normally using model = NeuralNetwork(). By default, its parameters are on the CPU, but we can move the model to a CUDA GPU using model = model.to("cuda").
But moving the model isn't enough, if our training data is still on the CPU then we have a problem.
for features, labels in train_loader:
predictions = model(features)Model → GPU
Features → CPU
↓
Can't perform the operation togetherSo we need to move the input tensors as well:
for features, labels in train_loader:
features = features.to("cuda")
labels = labels.to("cuda")
predictions = model(features)For example
Your code might run on a machine without a CUDA-compatible GPU. So you could choose the device first like so:
device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)And then use that throughout the application:
model = model.to(device)
for features, labels in train_loader:
features = features.to(device)
labels = labels.to(device)
predictions = model(features)Depending on the hardware you're targeting, PyTorch can also support other accelerator backends. For example, Apple Silicon Macs can use the mps device.
The general idea is still the same in that you simply decide which device you're using and make sure the relevant model parameters and tensors are there. (You can check where a tensor currently lives using print(features.device)).
And if you need to move a tensor back to the CPU you can use features = features.to("cpu"), or more commonly features = features.cpu().
This comes up quite often when you want to use the result with something outside PyTorch.
For example
If you want to convert a GPU tensor into a NumPy array, then you first need to move it back to the CPU with array = tensor.cpu().numpy().
One other useful detail is that calling .to(device) doesn't necessarily modify the original tensor in place, so features.to(device) isn't something you should rely on to update features. Instead, assign the result using features = features.to(device).
For an nn.Module, you'll commonly see model.to(device) because Module.to() moves its parameters and buffers in place and also returns the module.
So in an interview, you could explain it like this:
"I'd choose the device my application should use, move the model there with model.to(device), and move each batch of input tensors to the same device with tensor.to(device).
The model and tensors participating in an operation need to be on compatible devices. I'd usually detect whether the accelerator I want is available rather than hard-coding the device, and move tensors back to the CPU when necessary for CPU-only operations such as converting them to NumPy."
The important thing isn't simply knowing .to("cuda").
It's understanding that device placement is part of managing the entire computation:
Model parameters ─┐
├→ Same device → Computation
Input tensors ────┘Once you've trained that model, of course, you need some way to keep it rather than retraining it every time your application starts.
That brings us nicely to our next question…
#13. How do you save and load a PyTorch model?
Once you've spent time training a model, you obviously don't want to lose everything when your Python program stops running.
So you need to save what the model has learned so you can load it again later.
In PyTorch, the most common approach is to save the model's state_dict.
But remember that an nn.Module keeps track of things like its trainable parameters and registered buffers. Its state_dict is essentially a dictionary containing that state.
For example
Let's say that you run print(model.state_dict()). You might see entries representing the weights and biases of different layers like so:
layer1.weight
layer1.bias
layer2.weight
Layer2.biastorch.save(
model.state_dict(),
"model_weights.pth"
)Then, when you want to use the model again, you first create an instance of the same model architecture model = NeuralNetwork(), and load the saved state into it with:
model.load_state_dict(
torch.load("model_weights.pth", weights_only=True)
)So the basic process is:
Train model
↓
model.state_dict()
↓
torch.save()
↓
Saved weightsand later:
Create model architecture
↓
torch.load()
↓
model.load_state_dict()
↓
Restored modelOne important thing to understand though is that the state_dict doesn't define the model architecture for you. So, if you save the weights from our NeuralNetwork, you still need the code that defines NeuralNetwork when you load those weights again.
That's why the typical pattern is:
model = NeuralNetwork()
state_dict = torch.load(
"model_weights.pth",
weights_only=True
)
model.load_state_dict(state_dict)
model.eval()We call model.eval() here because if we're loading the model for inference, we want layers such as Dropout and BatchNorm to use their evaluation behavior, just like we discussed earlier.
You may also need to think about devices when loading a model.
For example
If you want to load the weights onto the CPU, you can use map_location:
state_dict = torch.load(
"model_weights.pth",
map_location="cpu",
weights_only=True
)
model.load_state_dict(state_dict)For example
Imagine training takes several days and your program crashes halfway through.
Sucks right?
Especially because if you've only saved the model parameters, you may not have everything you need to resume training exactly where you left off.
In that situation, you'd usually save a checkpoint containing additional training state:
torch.save({
"epoch": epoch,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"loss": loss
}, "checkpoint.pth")That's important for optimizers such as Adam, because the optimizer itself maintains information across training steps. Simply restoring the model weights isn't necessarily the same as restoring the entire training session.
So in an interview, you could explain it like this:
"The usual PyTorch approach is to save a model's state_dict with torch.save() and then recreate the model architecture and restore those parameters with load_state_dict().
If I need to resume training rather than just run inference, I'd save a checkpoint containing additional state such as the optimizer state and current epoch. I'd also put the loaded model into evaluation mode if I'm using it for inference."
The key distinction is:
Saving for inference
↓
Model state
Saving to resume training
↓
Model state
+ optimizer state
+ training progress
+ any other state you needThat brings us to our next question…
#14. How would you fine-tune a pretrained model in PyTorch?
You don't need to start with randomly initialized weights. Instead, you can take a model that's already been trained on a large dataset and adapt it to your own task. This is known as fine-tuning.
For example
Imagine you want to build an image classifier that recognizes five different types of flowers.
You could create a neural network from scratch and train it entirely on your flower dataset.
Or you could start with an image model that's already been trained on a large dataset and has learned useful visual features such as edges, shapes, textures, and more complex patterns.
Conceptually:
Pretrained model
↓
Already learned useful features
↓
Adapt it to your flower dataset
↓
Fine-tuned modelThe exact code depends on the model you're using, but let's imagine we've loaded a pretrained model.
One option is to fine-tune all of its parameters like so:
for param in model.parameters():
param.requires_grad = TrueFor example
If our image classifier was pretrained on ImageNet-1K, a subset of containing over 1 million training images across 1,000 classes, its output layer would originally predict those 1,000 classes.
But we only need it to distinguish between five flower classes, so we might replace its final layer:
import torch.nn as nn
model.classifier = nn.Linear(
model.classifier.in_features,
5
)Another approach is to freeze most of the pretrained model and train only the new classifier.
How?
Well, remember requires_grad from question #5?
This is one of the places where it becomes useful:
for param in model.parameters():
param.requires_grad = Falsemodel.classifier = nn.Linear(
model.classifier.in_features,
5
)So now we're effectively doing:
Pretrained layers
↓
Frozen
↓
Keep existing learned features
↓
New classifier
↓
TrainYou could also start this way and later unfreeze some or all of the pretrained layers for further fine-tuning.
The important thing is that fine-tuning isn't simply "load a pretrained model and train it." You need to decide which parameters should change.
If your dataset is relatively small or similar to the data the model originally learned from, freezing more of the pretrained model can help preserve useful features and reduce the amount of training required.
If your task is substantially different, or you have enough suitable training data, you might benefit from updating more of the model.
Learning rate matters too.
The pretrained parameters already contain useful information, so you often don't want to make huge updates that immediately destroy what the model has learned. It's common to use a smaller learning rate when fine-tuning pretrained parameters, and sometimes different learning rates for different parts of the model.
For example
optimizer = torch.optim.Adam([
{"params": model.features.parameters(), "lr": 1e-5},
{"params": model.classifier.parameters(), "lr": 1e-3}
])
So in an interview, you could explain it like this:
"I'd start with a pretrained model, adapt its output layer to my task, and decide which parameters I actually want to train. I might freeze the pretrained layers initially and train only the new head, or fine-tune some or all of the existing layers using a smaller learning rate. The choice depends on things like how much data I have and how similar my task is to the model's original training task."
The main idea is that fine-tuning lets you reuse useful representations the model has already learned rather than forcing your model to learn everything again from scratch.
But once you start training deeper networks, you can run into another problem: sometimes the gradients themselves become difficult to work with.
That brings us to our next question…
#15. What causes vanishing and exploding gradients, and how would you deal with them?
When training a neural network, gradients tell the optimizer how the model's parameters should change to reduce the loss. But sometimes those gradients become extremely small or extremely large as they're propagated backwards through the network.
These are known as vanishing and exploding gradients.
Let's start with vanishing gradients.
During backpropagation, gradients are calculated backwards through the layers of the network using the chain rule.
In a deep network, that means repeatedly multiplying derivatives together:
Output
↓
Gradient through layer 4
↓
× gradient through layer 3
↓
× gradient through layer 2
↓
× gradient through layer 1If many of those values are smaller than 1, repeatedly multiplying them can make the resulting gradient extremely small.
For example
0.1 × 0.1 × 0.1 × 0.1 = 0.0001By the time the gradient reaches earlier layers, it may be so small that their parameters barely change, and those layers effectively stop learning.
Exploding gradients are the opposite problem.
If the values being multiplied are large, gradients can grow rapidly:
5 × 5 × 5 × 5 = 625Now the optimizer may make enormous parameter updates, causing training to become unstable.
However, in practice, you might notice symptoms such as a loss that suddenly shoots upwards, parameters or gradients becoming extremely large, or even values turning into NaN or infinity.
For example
If I suspected a gradient problem in PyTorch, one useful first step would be to actually inspect the gradients.
for name, param in model.named_parameters():
if param.grad is not None:
print(name, param.grad.norm())What you do next depends on the cause.
One common technique for exploding gradients is gradient clipping.
PyTorch provides utilities such as:
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0
)optimizer.zero_grad()
predictions = model(features)
loss = loss_fn(predictions, labels)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0
)
optimizer.step()This prevents the overall gradient norm from exceeding the threshold you've chosen, which can help stabilize training, but gradient clipping isn't a universal fix.
I'd also look at things such as the learning rate, weight initialization, activation functions, normalization, and the architecture itself.
For example
Activation functions such as sigmoid can contribute to vanishing gradients in deep networks because their derivatives can become very small in saturated regions. ReLU and related activations often behave better in this respect.
Modern architectures also use techniques such as residual connections, which give gradients easier paths through deep networks.
And if the gradients are exploding because your learning rate is simply far too high, clipping them without addressing that underlying problem isn't necessarily the best solution.
So in an interview, you could explain it like this:
"Vanishing gradients happen when gradients become extremely small as they're propagated backwards through a network, so earlier layers learn very slowly or effectively stop learning. Exploding gradients are the opposite: the gradients become very large and training can become unstable. In PyTorch, I'd inspect gradient norms to confirm what's happening, and then look at the underlying cause.
Depending on the problem, solutions might include better initialization, suitable activation functions, normalization or residual connections, adjusting the learning rate, and gradient clipping for exploding gradients."
The important part of that answer is the diagnosis.
You don't want to hear "the loss is NaN " and immediately start randomly changing the model. You want to inspect what's happening during training and work out where things are going wrong.
Which actually gives us a very natural final question for the intermediate section…
#16. How would you debug a PyTorch model that isn't learning?
Sometimes your training loop runs perfectly, there are no errors, the GPU is working, and batches are happily moving through the model.
There's just one problem though and it's the fact that the model isn't getting any better.
Maybe the loss stays almost exactly the same:
Epoch 1: loss = 2.31
Epoch 2: loss = 2.30
Epoch 3: loss = 2.31
Epoch 4: loss = 2.30Or perhaps the loss decreases, but your validation performance doesn't improve.
When that happens, I wouldn't immediately start changing the architecture or trying random hyperparameters. I'd debug the training pipeline systematically.
How?
Well, the first thing I'd check is the data.
Are the inputs and labels actually what I think they are?
For example
features, labels = next(iter(train_loader))
print(features.shape)
print(labels.shape)
print(labels[:10])
This sounds basic, but it's surprisingly easy to spend hours debugging a model when the real problem is incorrect labels, unexpected preprocessing, the wrong tensor shape, or badly normalized input data.
Next, I'd check whether the forward pass and loss make sense.
For example
predictions = model(features)
print(predictions.shape)
loss = loss_fn(predictions, labels)
print(loss.item())Is the loss function appropriate for the task and compatible with those outputs?
For example
nn.CrossEntropyLoss expects raw logits for a classification problem, so applying a softmax yourself before passing the output to it would usually be unnecessary.
Then I'd check the gradients using loss.backward().
I can then inspect whether the trainable parameters are actually receiving gradients:
for name, param in model.named_parameters():
if param.grad is None:
print(name, "No gradient")
else:
print(name, param.grad.norm())If a parameter that should be learning has no gradient, I'd investigate why. Maybe it has requires_grad=False.
Maybe something in the computation has accidentally been detached from the autograd graph. Or perhaps the model contains a path where that parameter simply isn't contributing to the loss.
If gradients exist but are extremely small or huge, that could point toward the vanishing or exploding gradient problems we just discussed, so I'd also make sure the optimizer is actually updating the parameters I expect.
For example
If I created the optimizer before replacing a layer during fine-tuning, the new layer's parameters might not be in the optimizer at all. The code could run without an obvious error while the layer I care about never gets updated.
Then there are the training settings themselves.
A learning rate that's too small might produce almost no visible progress:
Weights barely change
↓
Loss barely changes
↓
Training appears stuckOne particularly useful debugging technique is to deliberately try to overfit a very small amount of data.
For example
Take just a handful of training examples and see whether the model can drive the training loss very low.
Tiny training set
↓
Can the model memorize it?
↙ ↘
Yes No
↓ ↓
Pipeline may Something may be
be working wrong with model,
loss, gradients,
optimizer, etc.If a model with enough capacity can't even learn a tiny training set, that's a strong signal that something is wrong with the training setup.
Finally, I'd distinguish between a model that isn't learning and one that isn't generalizing.
If training loss is decreasing nicely but validation performance is getting worse, the training loop may be working perfectly well, but the model could simply be overfitting.
That's a different problem and would lead me toward things such as regularization, data augmentation, model capacity, or additional training data rather than debugging gradient flow.
So in an interview, you could explain it like this:
"I'd debug a model that isn't learning systematically rather than immediately changing the architecture. I'd verify the data and labels, check tensor and output shapes, make sure the loss function is appropriate, inspect gradients, confirm the optimizer is updating the intended parameters, and check hyperparameters such as the learning rate.
I'd also try overfitting a tiny dataset as a sanity check. And I'd distinguish between failure to learn at all and overfitting, because those point to different problems."
The important thing is to work through the training pipeline and find evidence for what's actually going wrong.
A model that "isn't learning" is the symptom, but your job is to find the cause.
Advanced Pytorch interview questions
#17. How would you reduce GPU memory usage when training a PyTorch model?
As models and datasets get larger, one of the first practical limitations you can run into is GPU memory.
Why?
Well, you might have a perfectly valid training loop, but as soon as you increase the batch size or use a larger model, training crashes with an out-of-memory error.
So what is actually using all that memory?
Well during training, the GPU needs to hold more than just the model itself.
So at a simplified level, memory might be used for:
Model parameters
+
Gradients
+
Optimizer state
+
Activations from the forward pass
+
Current batch of dataRemember that during training, autograd needs information from the forward pass so it can calculate gradients during the backward pass. That means intermediate values may need to remain in memory until they're no longer required.
So if I were running out of GPU memory, I'd first try to understand what's consuming it rather than immediately changing the model.
One of the simplest options is reducing the batch size.
For example
train_loader = DataLoader(
dataset,
batch_size=64
)train_loader = DataLoader(
dataset,
batch_size=16
)But perhaps you wanted a larger effective batch size for training?
Well, in that case, you can use gradient accumulation. Instead of updating the model after every batch, you accumulate gradients across several smaller batches:
accumulation_steps = 4
optimizer.zero_grad()
for step, (features, labels) in enumerate(train_loader):
predictions = model(features)
loss = loss_fn(predictions, labels)
loss = loss / accumulation_steps
loss.backward()
if (step + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()Now four smaller batches can contribute to one optimizer update.
So if each batch contains 16 examples:
Batch 1: 16 ─┐
Batch 2: 16 │
Batch 3: 16 ├→ Accumulate gradients → Update
Batch 4: 16 ─┘
Effective batch size ≈ 64Another major technique is mixed precision training.
Instead of performing everything using 32-bit floating-point values, some operations can use lower-precision formats such as 16-bit floating point. This can significantly reduce memory use and can also speed up training on compatible hardware.
PyTorch provides automatic mixed precision, or AMP, to help manage this. (We'll cover AMP properly in the next question, because it's important enough to deserve its own answer).
For very large models though, another option is activation checkpointing, also known as gradient checkpointing.
Normally, autograd saves intermediate activations during the forward pass so they're available during backpropagation. But checkpointing trades computation for memory, instead of keeping all those intermediate activations:
Forward pass
↓
Save lots of activations
↓
Backward passForward pass
↓
Save fewer activations
↓
Backward pass
↓
Recompute missing values as neededTraining becomes slower because some computation happens twice, but memory usage can be substantially reduced.
PyTorch provides checkpointing utilities through torch.utils.checkpoint .
I'd also make sure I'm not accidentally retaining tensors or computational graphs that I no longer need.
For example
losses.append(loss) can cause trouble because loss is a tensor connected to the current computational graph. Keeping references to tensors like this can unintentionally keep associated data alive.
If I only want the numerical loss for logging, I'd normally use losses.append(loss.item()) .
Now I'm storing a Python number rather than retaining the tensor.
And for workloads that are simply too large to fit on one GPU even after these optimizations, the solution may need to become architectural. You might need to shard model parameters, gradients, and optimizer state across multiple GPUs using distributed approaches such as FSDP.
We'll get to distributed training shortly.
So in an interview, you could explain it like this:
"I'd first identify what's consuming the GPU memory. Common ways to reduce memory usage include lowering the batch size, using gradient accumulation if I still need a larger effective batch size, using automatic mixed precision, and using activation checkpointing to trade extra computation for lower activation memory.
I'd also check that I'm not accidentally retaining tensors or computation graphs. For models that still don't fit, I'd consider distributed approaches that shard model state across GPUs."
The important thing is that these techniques involve trade-offs.
Smaller batches
↓
Less memory
Gradient accumulation
↓
Larger effective batch without holding it all at once
Mixed precision
↓
Less memory + potentially faster
Activation checkpointing
↓
Less memory + more computation
Distributed sharding
↓
Spread memory requirements across multiple GPUsYou ask which memory-saving technique makes sense for the particular constraint you're dealing with.
And since mixed precision is one of the most useful of those techniques, let's look at it properly next.
m#18. What is automatic mixed precision (AMP), and why would you use it?
Neural networks perform huge numbers of numerical calculations during training.
Traditionally, many of those calculations use 32-bit floating-point numbers, often referred to as FP32.
At a simplified level:
FP32
↓
32 bits per value
↓
Good numerical range and precision
↓
More memory and computationBut not every operation during neural network training actually needs that level of precision.
If some operations can safely use lower-precision formats, we may be able to perform them faster and use less GPU memory.
That's the idea behind mixed precision training.
Instead of performing every operation using the same numerical precision, we use different precisions where appropriate.
For example
Some operations → lower precision
Some operations → FP32
The problem is that deciding which operations should use which precision yourself would be tedious and easy to get wrong.
That's where PyTorch's automatic mixed precision, or AMP, comes in. AMP can automatically choose an appropriate precision for operations during the forward pass.
On a CUDA GPU, a training loop might look something like:
scaler = torch.amp.GradScaler("cuda")
for features, labels in train_loader:
features = features.to("cuda")
labels = labels.to("cuda")
optimizer.zero_grad()
with torch.autocast(
device_type="cuda",
dtype=torch.float16
):
predictions = model(features)
loss = loss_fn(predictions, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()There are two important ideas happening here.
First, autocast controls the precision used during the forward pass:
with torch.autocast(
device_type="cuda",
dtype=torch.float16
):This can reduce memory usage and, on supported hardware, significantly improve performance.
However, using FP16 introduces another problem.
What?
Very small gradient values can become too small to represent properly in FP16. This is known as underflow.
For example
Imagine a gradient becoming smaller and smaller:
0.00001
↓
0.000001
↓
Too small for the available precision
↓
Effectively becomes zeroIf that happens repeatedly, useful gradient information can disappear. That's why FP16 training commonly uses gradient scaling.
So instead of immediately backpropagating the original loss with loss.backward() , we use scaler.scale(loss).backward() . The scaler temporarily multiplies the loss, which makes the resulting gradients larger and less likely to underflow.
Before the optimizer actually updates the parameters, PyTorch accounts for that scaling appropriately. GradScaler can also adjust the scale dynamically if it detects numerical problems.
So the overall process looks roughly like:
Forward pass
↓
Use mixed precision
↓
Calculate loss
↓
Scale loss
↓
Backward pass
↓
Safely update parametersBF16 has a larger numerical range than FP16, which means gradient scaling isn't generally required in the same way. The best precision therefore depends on your hardware and workload, and AMP isn't guaranteed to make every model faster.
The actual benefit depends on things like the operations your model performs and the hardware you're running on. So just like our other production optimizations, I'd measure the effect rather than assuming it will automatically help.
So in an interview, you could explain it like this:
"Automatic mixed precision lets PyTorch run suitable operations at lower precision while keeping other operations at a precision that's safer for numerical stability. This can reduce GPU memory usage and improve training speed on supported hardware.
With FP16, I'd typically use autocast for the forward pass and GradScaler to help prevent small gradients from underflowing during backpropagation. I'd also benchmark it for the particular model and hardware rather than assuming it will always be faster."
The key idea is that you're not simply converting the entire model to FP16 and hoping for the best. AMP lets you get many of the performance and memory benefits of lower precision while PyTorch helps manage where that lower precision is appropriate.
So far, though, we've still been assuming that our training happens on a single device.
For larger workloads, you may want to spread the work across several GPUs or even several machines.
#19. How does distributed training work in PyTorch?
So far, we've assumed that our model is being trained on a single GPU, but eventually you might have a workload where training on one GPU takes too long, or the model itself is too large to fit comfortably on a single device.
That's where distributed training comes in.
At a high level, distributed training means spreading the training workload across multiple GPUs, and potentially multiple machines.
There are different ways to do that, but one of the most common is data parallelism.
For example
Imagine we have four GPUs, instead of sending an entire batch through one GPU:
Batch of 256
↓
GPU 0
↓
Forward + backward passWe can give each GPU part of the batch:
Batch of 256
↓
┌───────────┼───────────┐
↓ ↓ ↓ ↓
GPU 0 GPU 1 GPU 2 GPU 3
64 samples 64 samples 64 samples 64 samples
↓ ↓ ↓ ↓
Local Local Local Local
gradients gradients gradients gradients
└───────────┼───────────┘
↓
Synchronize gradients
↓
Update parametersIf every GPU independently updated its model using only its own gradients, we'd quickly end up with four different models. So instead, the gradients are synchronized across the processes so that each model replica makes a consistent parameter update.
In PyTorch, the standard approach for this type of training is DistributedDataParallel , usually shortened to DDP.
With DDP, you'd typically run one process per GPU.
Conceptually:
Process 0 → GPU 0 → Model replica
Process 1 → GPU 1 → Model replica
Process 2 → GPU 2 → Model replica
Process 3 → GPU 3 → Model replicaEach process handles its own portion of the training data.
That's another important detail though because you don't want every GPU repeatedly training on exactly the same examples.
A distributed sampler can partition the dataset so each process receives a different portion:
from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(
train_dataset,
batch_size=64,
sampler=sampler
)The model can then be wrapped with DDP like so:
from torch.nn.parallel import DistributedDataParallel as DDP
model = model.to(local_rank)
model = DDP(
model,
device_ids=[local_rank]
)
The important thing is understanding the architecture.
DDP is useful when the model itself fits on a single GPU, but you want to speed up training by processing data across multiple GPUs.
However, adding GPUs doesn't automatically make training proportionally faster.
The processes now need to communicate with each other to synchronize gradients. That communication introduces overhead, and the benefit depends on things like model size, batch size, hardware, and how quickly the GPUs can communicate.
There's also a different problem distributed training can solve.
What?
Well, what if the model itself is too large to fit on one GPU? In that case, making four copies of it with DDP obviously doesn't help:
Model requires 40 GB
GPU 0 has 24 GB → Doesn't fit
GPU 1 has 24 GB → Doesn't fit
GPU 2 has 24 GB → Doesn't fit
GPU 3 has 24 GB → Doesn't fitSo in that situation, you need a strategy that distributes the model's memory requirements rather than replicating the whole model on every GPU.
PyTorch provides approaches such as Fully Sharded Data Parallel, or FSDP, which can shard model parameters, gradients, and optimizer state across workers.
Conceptually:
One large model
↓
Shard its state
↓
┌─────┼─────┬─────┐
GPU 0 GPU 1 GPU 2 GPU 3
Model fits on one GPU,
but training is too slow
↓
Data parallelism / DDPModel itself doesn't fit
on one GPU
↓
Sharding approaches such as FSDPSo in an interview, you could explain it like this:
"Distributed training spreads training work across multiple devices or machines. With DistributedDataParallel, I'd normally run one process per GPU, give each process a different portion of the data, and let each GPU perform its own forward and backward passes. DDP then synchronizes gradients so the model replicas remain consistent.
If the model itself is too large to fit on a single GPU, I'd look at a sharded approach such as FSDP rather than simply replicating the entire model."
The important thing is to understand what you're trying to distribute. Sometimes you're distributing the data processing to make training faster. Other times you're distributing the model state because it won't fit on one device.
And there's another PyTorch API you may encounter when talking about multi-GPU training: DataParallel.
It sounds very similar to DistributedDataParallel, but they're not the same thing which brings us to our next question…
#20. What's the difference between DataParallel and DistributedDataParallel?
PyTorch has two APIs whose names make them sound like they do roughly the same thing. torch.nn.DataParallel and torch.nn.parallel.DistributedDataParallel.
Both can be used to train a model across multiple GPUs, but if you're building a modern multi-GPU training setup, DistributedDataParallel, or DDP, is generally the approach you should prefer.
Let's look at why.
DataParallel uses a single process to coordinate multiple GPUs.
You give it a model:
model = torch.nn.DataParallel(model)
model = model.to("cuda")And during the forward pass it splits the input batch across the available GPUs.
Conceptually:
One process
↓
Split the batch
↓
┌───────────┼───────────┐
↓ ↓ ↓
GPU 0 GPU 1 GPU 2
↓ ↓ ↓
Model Model Model
replica replica replica
└───────────┼───────────┘
↓
Gather the resultsThis makes DataParallel relatively simple to use. The problem is that the single process becomes a bottleneck. This is because one GPU also takes on additional coordination work, including gathering the outputs, so the workload isn't distributed as efficiently as it could be.
DDP uses a different architecture.
Instead of one process controlling all the GPUs, you typically run one process for each GPU:
Process 0 → GPU 0
Process 1 → GPU 1
Process 2 → GPU 2
Process 3 → GPU 3The GPUs perform their forward and backward passes independently, and DDP synchronizes the gradients between them during backpropagation.
As we saw in the previous question:
GPU 0 ─→ gradients ─┐
GPU 1 ─→ gradients ─┤
GPU 2 ─→ gradients ─┼→ Synchronize
GPU 3 ─→ gradients ─┘
↓
Consistent updatesThis architecture avoids the same single-process bottleneck and generally scales much better.
It can also work across multiple machines, whereas DataParallel is designed for multiple GPUs on a single machine.
So although both approaches replicate the model and split work across GPUs, their architecture is quite different:
DataParallel
One process
↓
Multiple GPUs
versus:
DistributedDataParallel
Multiple processes
↓
Usually one process per GPUBoth generally involve having a model replica on each GPU, so if your problem is "My model is too large to fit on one GPU", switching from DataParallel to DDP doesn't magically solve it.
That's when approaches that shard the model state, such as the FSDP approach we mentioned earlier, become relevant.
So in an interview, you could explain it like this:
"DataParallel uses a single process to split batches across multiple GPUs, while DistributedDataParallel normally uses one process per GPU and synchronizes gradients between those processes. DDP avoids many of the bottlenecks of DataParallel, scales better, and can also work across multiple machines, so it's generally the preferred approach for multi-GPU training.
Neither is primarily a solution for a model that's too large to fit on one GPU, because both use model replicas."
The most important thing isn't memorizing that DDP is "better."
It's understanding why:
DataParallel
↓
Simple, but centralized
single-process coordination
DDP
↓
Separate processes
+ efficient gradient synchronization
+ better scaling#21. What is torch.compile, and when would you use it?
Normally, PyTorch executes your model operations as Python runs them. This easy execution is one of the things that makes PyTorch pleasant to work with, because you can write ordinary Python, run the code immediately, and easily inspect what's happening.
Handy right?
However, executing operations individually isn't always the fastest way to run a model and that's where torch.compile comes in.
Introduced in PyTorch 2, torch.compile can optimize a model or function by capturing regions of its computation and compiling them into more efficient code.
It's fairly simple to use also:
model = NeuralNetwork()
compiled_model = torch.compile(model)predictions = compiled_model(features)Well, imagine part of your model performs several operations:
Input
↓
Operation A
↓
Operation B
↓
Operation C
↓
Operation D
↓
OutputWith normal eager execution, PyTorch executes those operations as the Python program encounters them. But with torch.compile, PyTorch can capture the computation and optimize how groups of operations are executed.
At a simplified level:
PyTorch model
↓
torch.compile()
↓
Capture computation
↓
Optimize it
↓
Generate efficient code
↓
Run compiled computationOne potential optimization is reducing the overhead of launching lots of separate operations.
For example
Several compatible operations may be combined into a smaller number of optimized kernels, a technique known as kernel fusion.
Instead of:
GPU kernel 1
↓
GPU kernel 2
↓
GPU kernel 3You may be able to do something closer to Optimized fused kernel. That can reduce overhead and unnecessary movement of data, potentially making the model faster.
The nice thing is that you don't usually have to rewrite your model in an entirely different framework to try it.
You can start with ordinary PyTorch code and compile it with model = torch.compile(model).
However, that doesn't mean every model automatically becomes faster, simply because compilation itself has a cost.
The first few executions can be slower because PyTorch needs to capture and compile the relevant computation before it can benefit from the optimized version. So if you only run a model once or twice, the compilation overhead may outweigh the performance improvement.
Another consideration is that some Python or PyTorch operations can cause a graph break.
A graph break means PyTorch can't include that part of the program in the current compiled graph, so execution falls back to normal Python before compilation potentially resumes later.
Conceptually:
Compiled region
↓
Unsupported/dynamic behavior
↓
Graph break
↓
Eager execution
↓
Compiled regionA few graph breaks aren't necessarily disastrous, but lots of them can reduce the benefit of compiling the model.
Dynamic behavior can also sometimes cause PyTorch to recompile when it encounters inputs or execution paths that don't match what it previously compiled.
So if torch.compile isn't giving me the performance improvement I expected, I'd want to investigate things such as compilation time, graph breaks, recompilations, and which parts of the workload are actually being optimized.
And most importantly, I'd benchmark it.
For a long-running training job where the same model executes thousands of times, paying an upfront compilation cost could easily make sense if every subsequent iteration becomes faster.
For a tiny workload that only runs a handful of times, it might not.
So in an interview, you could explain it like this:
"torch.compile is a PyTorch feature that can optimize models and functions by capturing their computations and compiling them into more efficient code. It can improve performance through optimizations such as kernel fusion while still letting you write normal PyTorch code.
But compilation has upfront overhead, and graph breaks or recompilations can reduce the benefit, so I'd benchmark it on the actual workload rather than assuming it will always make a model faster."
The key idea is that torch.compile isn't a magic "make my model fast" switch.
It's another optimization tool, and whether it helps depends on what is actually limiting your application.
#22. How would you identify a performance bottleneck in a PyTorch model?
We've looked at several ways to make PyTorch models faster, from mixed precision and distributed training to torch.compile.
But there's an important step that should usually come before any of those optimizations and that’s to work out what's actually slow. Otherwise, you can spend time optimizing something that wasn't limiting performance in the first place.
For example
Suppose training is taking much longer than expected.
Well, the bottleneck could be almost anywhere:
Loading and preprocessing data
↓
Moving data to the GPU
↓
Forward pass
↓
Backward pass
↓
Optimizer step
↓
Synchronizing between GPUsMaybe your GPU is waiting for the CPU to prepare each batch
Maybe you're constantly transferring data between the CPU and GPU
Or perhaps a particular operation inside the model is consuming most of the execution time
So rather than guessing, I'd profile the workload.
PyTorch provides torch.profiler specifically for collecting performance information during training and inference. It can measure CPU activity, accelerator activity, operator execution time, memory consumption, input shapes, and more.
For example
We could profile CPU and CUDA activity like this:
from torch.profiler import profile, ProfilerActivity
with profile(
activities=[
ProfilerActivity.CPU,
ProfilerActivity.CUDA
],
record_shapes=True,
profile_memory=True
) as prof:
predictions = model(features)
loss = loss_fn(predictions, labels)
loss.backward()
optimizer.step()Then we can inspect the results:
print(
prof.key_averages().table(
sort_by="cuda_time_total",
row_limit=10
)
)For example
Which operations take the most time?
Is most of the time spent on the CPU or GPU?
Are particular operations allocating lots of memory?
Is the GPU doing useful work or waiting for data?
Are CPU ↔ GPU transfers taking too long?
The profiler can also record tensor shapes and stack information, which can help you trace an expensive operation back to the code responsible for it. PyTorch can export profiler results as traces too, making it easier to inspect the timing and relationship between CPU operations and accelerator kernels.
Once you've identified the bottleneck, you can choose an optimization that actually addresses it.
For example
Data loading bottleneck
↓
Improve DataLoader configuration or preprocessing
GPU memory bottleneck
↓
Smaller batches, AMP, checkpointing, etc.
Expensive model operations
↓
Investigate those operations or try torch.compile
Multi-GPU communication bottleneck
↓
Investigate distributed training configurationFor example
Suppose you enable AMP because you've heard it makes training faster. That's useful if lower-precision computation improves your particular workload.
But the real process should be:
Measure baseline
↓
Identify bottleneck
↓
Make a change
↓
Measure again
↓
Did it actually improve?The same applies to torch.compile because you shouldn't assume torch.compile = faster.
You should benchmark the model before and after compilation and see whether the improvement outweighs things like compilation overhead.
There's one small distinction worth making here too because profiling and benchmarking aren't quite the same thing.
Profiling helps answer “Where is the time going?”
While benchmarking helps answer “How fast is this code?”
PyTorch also provides benchmarking utilities such as torch.utils.benchmark, but when you're trying to understand why an entire training pipeline is slow, torch.profiler gives you much more visibility into what's happening inside it.
PyTorch's current documentation describes the profiler specifically as a way to understand expensive model operators, input shapes, stack traces, device kernel activity, and execution traces.
So in an interview, you could explain it like this:
"I'd measure the workload before trying to optimize it. In PyTorch, I'd use torch.profiler to inspect CPU and GPU activity, operator execution times, memory usage, and potentially input shapes or stack traces. That helps determine whether the bottleneck is data loading, device transfers, model computation, memory pressure, or something else.
I'd then optimize the actual bottleneck and benchmark the workload again to verify that the change improved performance."
Wrapping up
And there you have it, 22 PyTorch interview questions and answers to help you prepare for your next interview.
How did you get on?
If you could explain most of these answers in your own words, rather than simply remembering the definitions, you're in a pretty good place
But if there were a few questions where you got stuck, then I'd spend some time brushing up on those areas before your interview. Better still, try building something with them
The best way to build that understanding is to actually use PyTorch. Build a few models, train them, break them, debug them, and experiment with the techniques we've covered here.
And if you want a structured way to do that, check out my PyTorch course:
You'll learn PyTorch by building real projects from scratch, so you're not just preparing to answer interview questions.
Either way, good luck with your interview!
You've got this.
Best articles. Best resources. Only for ZTM subscribers.
If you enjoyed this post and want to get more like it in the future, subscribe below. By joining the ZTM community of over 100,000 developers you’ll receive Web Developer Monthly (the fastest growing monthly newsletter for developers) and other exclusive ZTM posts, opportunities and offers.
No spam ever, unsubscribe anytime
