PyTorch Interview Questions for Machine Learning, with Code

What PyTorch questions are asked in ML interviews?

ML engineer and research engineer interviews ask PyTorch questions about tensors and their memory layout, autograd, nn.Module, losses, optimizers and schedulers, the training loop, train and eval mode, Dataset and DataLoader, devices, mixed precision, saving checkpoints and debugging shape and NaN problems. TensorTonic has a PyTorch Sheet of 30 problems in six sections for these topics, and this guide gives a short answer and a small code example for each common question, with a link to a practice problem where one exists.

The questions usually arrive as a short piece of code where you predict a shape, a gradient or an error, or as a task such as writing a custom layer or a full training loop, followed by a question about what would break. The conceptual deep learning questions, such as why batch norm works or how Adam differs from SGD, are in our deep learning interview questions guide, and the order to learn PyTorch in is in how to learn PyTorch. The code on this page assumes import torch and from torch import nn, and every snippet runs on a CPU with current PyTorch.

What do interviewers ask about PyTorch tensors?

Interviewers ask how tensors differ from NumPy arrays, when an operation shares memory instead of copying, and how broadcasting and in-place operations can change results without raising an error. These questions check that you understand what a tensor is underneath the API.

import numpy as np
import torch

a = np.array([1.0, 2.0, 3.0])
shared = torch.from_numpy(a)
copied = torch.tensor(a)
a[0] = 100.0
print(shared[0].item(), copied[0].item())
print(shared.dtype, torch.tensor([1.0, 2.0]).dtype)
w = torch.ones(3, requires_grad=True)
print(w.detach().numpy())

x = torch.arange(6).reshape(2, 3)
t = x.t()
print(x.stride(), t.stride(), t.is_contiguous())
try:
    t.view(6)
except RuntimeError as err:
    print(type(err).__name__)
flat = t.reshape(6)
also_flat = t.contiguous().view(6)
print(torch.equal(flat, also_flat))
v = x.view(3, 2)
v[0, 0] = 99
print(x[0, 0].item())

pred = torch.randn(4, 1)
target = torch.randn(4)
print((pred - target).shape)
print((pred.squeeze(1) - target).shape)

w = torch.randn(3, requires_grad=True)
y = w.exp()
y.add_(1)
try:
    y.sum().backward()
except RuntimeError as err:
    print(type(err).__name__)

with torch.no_grad():
    w.sub_(0.1 * torch.ones(3))
print(w.requires_grad, w.is_leaf)

What autograd questions do interviewers ask?

Interviewers ask what requires_grad records, why gradients accumulate, and how to stop autograd from tracking a computation. Autograd is easy to use without understanding it, so interviewers probe it with small examples where you predict the value of a grad attribute.

x = torch.tensor([2.0, 3.0], requires_grad=True)
h = x * x
h.retain_grad()
loss = (h * 4).sum()
loss.backward()
print(x.is_leaf, h.is_leaf, h.grad_fn is not None)
print(x.grad, h.grad)

model = nn.Linear(4, 1)
x = torch.randn(8, 4)
y = torch.randn(8, 1)

model.zero_grad()
nn.functional.mse_loss(model(x), y).backward()
full = model.weight.grad.clone()

model.zero_grad()
print(model.weight.grad)
for xb, yb in zip(x.split(2), y.split(2)):
    (nn.functional.mse_loss(model(xb), yb) / 4).backward()
print(torch.allclose(full, model.weight.grad))

w = torch.randn(3, requires_grad=True)
a = (w * 2).detach()
with torch.no_grad():
    b = w * 2
with torch.inference_mode():
    c = w * 2
print(a.requires_grad, b.requires_grad, c.requires_grad, c.is_inference())
try:
    (c * w).sum().backward()
except RuntimeError as err:
    print(type(err).__name__)
print(w.detach().data_ptr() == w.data_ptr())

What nn.Module questions do interviewers ask?

Interviewers ask how a module registers its parameters and buffers and how to write a custom layer from memory. A good answer shows that you know what ends up in parameters() and in the state_dict, because that decides what gets trained and what gets saved.

class Scaled(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(dim))
        self.register_buffer("running_mean", torch.zeros(dim))
        self.layers = nn.ModuleList([nn.Linear(dim, dim) for _ in range(2)])
        self.hidden = [nn.Linear(dim, dim)]

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return (x - self.running_mean) * self.scale


m = Scaled(4)
print(sum(p.numel() for p in m.parameters()))
print(list(m.state_dict().keys()))

class MyLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
        nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5)

    def forward(self, x):
        return x @ self.weight.T + self.bias


layer = MyLinear(3, 2)
x = torch.randn(5, 3)
out = layer(x)
out.sum().backward()
print(out.shape, layer.weight.grad.shape)
ref = nn.Linear(3, 2)
ref.load_state_dict(layer.state_dict())
print(torch.allclose(ref(x), out))

What does nn.CrossEntropyLoss expect as input?

The nn.CrossEntropyLoss class expects raw logits of shape (N, C) and integer class indices of shape (N,) with dtype int64, and it applies log-softmax itself, so it equals nll_loss on log_softmax output, as the first two lines below show. It also accepts class probabilities of shape (N, C) as the target for soft labels. For sequence models you flatten the batch and time dimensions together, and any target equal to ignore_index, which defaults to -100, is left out of the loss, which is the usual way to mask padding tokens. For binary or multi-label tasks use nn.BCEWithLogitsLoss, which takes logits and float targets of the same shape. The deep learning interview questions guide explains why losses take logits, and the guide on which loss function to use matches losses to tasks. Practise with Loss Functions, Softmax from Scratch and Cross-Entropy Loss.

logits = torch.tensor([[2.0, 0.5, -1.0], [0.1, 0.2, 3.0]])
labels = torch.tensor([0, 2])
print(nn.CrossEntropyLoss()(logits, labels))
print(nn.functional.nll_loss(logits.log_softmax(dim=1), labels))

seq_logits = torch.randn(2, 5, 10)
tokens = torch.randint(0, 10, (2, 5))
tokens[1, 3:] = -100
seq_loss = nn.functional.cross_entropy(seq_logits.reshape(-1, 10), tokens.reshape(-1), ignore_index=-100)

bce = nn.BCEWithLogitsLoss()
print(bce(torch.tensor([1.5, -0.3]), torch.tensor([1.0, 0.0])))

In what order do you call the optimizer and the scheduler?

You call optimizer.zero_grad(), then loss.backward(), then optimizer.step(), and only after that scheduler.step(), and PyTorch prints a warning if the scheduler steps before the optimizer has. Most schedulers for transformers, such as warmup followed by cosine decay, step once per batch, while older schedules such as StepLR are usually stepped once per epoch, so say which one you mean. Gradient clipping goes between backward and step, since it needs the gradients to exist and must change them before the update uses them. A LambdaLR multiplies the base learning rate by whatever your function returns for the current step, which makes warmup a one-line function. Practise with Optimizer Scheduler, Learning Rate Warmup Scheduler, Custom SGD with Momentum and Gradient Clipping (Global Norm), and read the guide on which optimizer to use for choosing between them.

model = nn.Linear(10, 1)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
warmup = 100
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda step: min(1.0, (step + 1) / warmup))

x, y = torch.randn(32, 10), torch.randn(32, 1)
for step in range(3):
    opt.zero_grad()
    loss = nn.functional.mse_loss(model(x), y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    opt.step()
    sched.step()
    print(step, sched.get_last_lr())

How do you write a PyTorch training loop in an interview?

You write a PyTorch training loop as one function that trains for an epoch and one that evaluates, where the training function moves each batch to the device, clears the gradients, runs the forward pass and the loss, calls backward and steps the optimizer. The evaluation function switches to eval mode, runs under torch.no_grad and returns the average loss. Weight each batch loss by the batch size before averaging, because the last batch is often smaller than the rest, and call loss.item() so the running total does not hold on to the graph. Interviewers often ask for this from a blank editor, and the how to learn PyTorch guide lists the same steps in order. Practise with Mini Training Loop and add validation-based stopping with Early Stopping.

from torch.utils.data import DataLoader, TensorDataset


def train_one_epoch(model, loader, loss_fn, opt, device):
    model.train()
    total = 0.0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        opt.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        opt.step()
        total += loss.item() * xb.size(0)
    return total / len(loader.dataset)


@torch.no_grad()
def evaluate(model, loader, loss_fn, device):
    model.eval()
    total = 0.0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        total += loss_fn(model(xb), yb).item() * xb.size(0)
    return total / len(loader.dataset)


device = "cuda" if torch.cuda.is_available() else "cpu"
X, y = torch.randn(256, 8), torch.randint(0, 3, (256,))
train_loader = DataLoader(TensorDataset(X[:200], y[:200]), batch_size=32, shuffle=True)
val_loader = DataLoader(TensorDataset(X[200:], y[200:]), batch_size=64)
model = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 3)).to(device)
opt = torch.optim.Adam(model.parameters(), lr=1e-2)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(3):
    train_loss = train_one_epoch(model, train_loader, loss_fn, opt, device)
    val_loss = evaluate(model, val_loader, loss_fn, device)
    print(epoch, round(train_loss, 3), round(val_loss, 3))

What do model.train() and model.eval() change?

The model.train() and model.eval() methods set the training flag on every submodule, which changes how layers such as dropout and batch norm behave and nothing else. In training mode dropout zeroes random activations and batch norm normalises with the statistics of the current batch and updates its running averages, and in eval mode dropout passes its input through and batch norm uses the running averages. Neither method turns off gradient tracking, which is why the output below still has requires_grad set, so evaluation code needs torch.no_grad or torch.inference_mode as well. A model that scores far worse on its own training data in eval mode than in train mode often has batch norm running averages that do not match the data. Practise both layers with Batch Normalization and Dropout from Scratch.

net = nn.Sequential(nn.Linear(4, 4), nn.BatchNorm1d(4), nn.Dropout(0.5))
x = torch.randn(8, 4)
net.train()
print(torch.equal(net(x), net(x)))
net.eval()
print(torch.equal(net(x), net(x)))
print(net(x).requires_grad)
print(net[1].running_mean)

What is the difference between a Dataset and a DataLoader?

A Dataset returns one example at a time through __getitem__ and reports how many examples it has through __len__, while a DataLoader wraps a Dataset and turns it into batches, with options for shuffling, the number of worker processes and pinned memory. The DataLoader combines examples into a batch with a collate_fn, and the default one stacks tensors of the same shape, so variable-length sequences need a custom collate_fn that pads them, as in the example. Setting num_workers above zero loads batches in background processes so the GPU does not wait for data, and pin_memory=True puts batches in page-locked memory so that copying them to the GPU with non_blocking=True can overlap with computation. Shuffle the training loader and leave the validation loader in order. Practise with Custom Dataset Class, Balanced DataLoader and Transform Pipeline.

from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset


class TokenDataset(Dataset):
    def __init__(self, sequences, labels):
        self.sequences = sequences
        self.labels = labels

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        return torch.tensor(self.sequences[idx]), torch.tensor(self.labels[idx])


def collate(batch):
    seqs, labels = zip(*batch)
    padded = pad_sequence(seqs, batch_first=True, padding_value=0)
    return padded, torch.stack(labels)


ds = TokenDataset([[5, 2, 9], [7, 1], [3, 3, 3, 3]], [1, 0, 1])
loader = DataLoader(ds, batch_size=2, shuffle=True, collate_fn=collate)
for xb, yb in loader:
    print(xb.shape, yb)

How do you move a model and its data to a GPU?

You move a model to a GPU with model.to(device) and each batch with xb = xb.to(device), and every tensor in an operation must be on the same device or PyTorch raises an error that names both devices. The to method moves a module's parameters and buffers in place, but on a tensor it returns a new tensor and leaves the original where it was, so the unassigned call in the example does nothing. Create new tensors inside forward with device=x.device so the code runs wherever its input lives, and move the model before you create the optimizer so the optimizer holds the moved parameters. Bring results back with cpu() before converting them to NumPy or Python values. On a Mac the device string is mps instead of cuda. Our guide on how to practise ML without a GPU covers what you can do on a CPU alone.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = nn.Linear(4, 2).to(device)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
x = torch.randn(3, 4)
x.to(device)
x = x.to(device)
mask = torch.ones(3, 2, device=x.device)
out = model(x) * mask
print(out.device, next(model.parameters()).device)
preds = out.argmax(dim=1).cpu().tolist()

How does mixed precision training work with autocast?

Mixed precision training runs selected operations such as matrix multiplications in float16 or bfloat16 inside a torch.autocast block, while the weights stay in float32 and operations that need precision, such as the loss, run in float32. With float16 you also need a GradScaler, which multiplies the loss by a large factor before backward so that small gradients do not round to zero, then divides the gradients again before the optimizer step and skips steps whose gradients overflowed. With bfloat16 you usually do not need a scaler, because it has the same exponent range as float32. Call scaler.unscale_ before clipping gradients, since clipping scaled gradients would use the wrong threshold. The example runs float16 with a scaler on a GPU and bfloat16 without one on a CPU. Read mixed precision training for how the formats differ, and practise the memory side with Mixed-Precision Training Memory Accountant.

device = "cuda" if torch.cuda.is_available() else "cpu"
use_fp16 = device == "cuda"
dtype = torch.float16 if use_fp16 else torch.bfloat16
model = nn.Linear(16, 4).to(device)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
scaler = torch.amp.GradScaler(device, enabled=use_fp16)

x = torch.randn(8, 16, device=device)
y = torch.randint(0, 4, (8,), device=device)
with torch.autocast(device_type=device, dtype=dtype):
    logits = model(x)
    loss = nn.functional.cross_entropy(logits, y)
print(logits.dtype, loss.dtype, model.weight.dtype)

opt.zero_grad()
scaler.scale(loss).backward()
scaler.unscale_(opt)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(opt)
scaler.update()

How should you save and load a PyTorch model?

You should save a model's state_dict with torch.save and load it into a freshly built model with load_state_dict, instead of pickling the whole model object, which ties the file to your exact class definitions and file layout. A checkpoint that can resume training also stores the optimizer's state_dict, the scheduler's state_dict if there is one, and the epoch or step count. Pass map_location when you load on a machine with a different device, and keep weights_only=True, which has been the default since PyTorch 2.6 and refuses to run arbitrary pickled code. Keys that start with module. come from saving a model wrapped in DataParallel or DistributedDataParallel, and keys that start with _orig_mod. come from a torch.compile wrapper, so strip the prefix or save the inner model's state_dict. The printed result of load_state_dict lists any missing or unexpected keys, which is the first thing to read when loading fails.

model = nn.Linear(4, 2)
opt = torch.optim.Adam(model.parameters())
torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "epoch": 5}, "checkpoint.pt")

ckpt = torch.load("checkpoint.pt", map_location="cpu", weights_only=True)
restored = nn.Linear(4, 2)
restored.load_state_dict(ckpt["model"])
restored_opt = torch.optim.Adam(restored.parameters())
restored_opt.load_state_dict(ckpt["optimizer"])
start_epoch = ckpt["epoch"] + 1

wrapped = {"module." + k: v for k, v in model.state_dict().items()}
clean = {k.removeprefix("module."): v for k, v in wrapped.items()}
print(restored.load_state_dict(clean))

How do you debug shape errors and NaN losses in PyTorch?

You debug a shape error by printing the shape after each layer until you find the first one that differs from what you expected, and the error text usually tells you where to look, as in the message that mat1 and mat2 shapes cannot be multiplied, which gives both shapes. You debug a NaN loss by checking whether the loss and the gradients are finite after every step, then looking for the usual causes, which are a learning rate that is too high, a log or division of zero, float16 overflow, and exploding gradients. Use log_softmax or a loss that takes logits instead of taking the log of a probability yourself, because the stable versions never see an exact zero. Turning on torch.autograd.set_detect_anomaly(True) makes backward report the forward operation that produced the first NaN, at the cost of much slower training. Our guide on a training loss that is not decreasing covers the case where the loss stays finite but does not fall, and Learning Rate Warmup Scheduler practises the warmup that prevents many early NaNs.

model = nn.Sequential(nn.Linear(10, 16), nn.ReLU(), nn.Linear(16, 1))
x = torch.randn(4, 10)
for name, module in model.named_children():
    x = module(x)
    print(name, tuple(x.shape))

try:
    model(torch.randn(4, 12))
except RuntimeError as err:
    print(err)


def check_finite(model, loss):
    if not torch.isfinite(loss):
        raise ValueError("loss is not finite")
    for name, p in model.named_parameters():
        if p.grad is not None and not torch.isfinite(p.grad).all():
            raise ValueError(f"gradient of {name} is not finite")


print(torch.log(torch.tensor([0.0, -1.0])))
print(nn.functional.binary_cross_entropy_with_logits(torch.tensor([200.0]), torch.tensor([0.0])))

What does torch.compile do?

The torch.compile function captures the Python operations of a model as a graph the first time it runs and hands that graph to a compiler, which by default is TorchInductor, so later calls run generated code that fuses many small operations into fewer kernels. On a GPU Inductor generates Triton kernels, and on a CPU it generates C++. The first call is slow because it compiles, and new input shapes or Python control flow that the graph cannot capture cause recompiles or graph breaks, which reduce the speedup. The compiled model shares its parameters with the original, and its state_dict keys carry the _orig_mod. prefix. Interviewers rarely go deeper than this unless the role is about performance, in which case the how to learn Triton guide and the Fused Tiled Matmul and ReLU problem show what kernel fusion means in practice.

model = nn.Sequential(nn.Linear(64, 64), nn.GELU(), nn.Linear(64, 64))
compiled = torch.compile(model)
x = torch.randn(32, 64)
print(torch.allclose(compiled(x), model(x), atol=1e-5))
print(next(iter(compiled.state_dict())))

How should you practise PyTorch interview questions?

You should practise PyTorch interview questions by writing the code for each one yourself and running it, starting with the TensorTonic PyTorch Sheet, which has 30 problems in six sections. The sections run from foundations and building blocks through the training loop and training infrastructure to advanced layers and a final section on sequence models, with tensors and autograd at the start and attention and beam search at the end. Every solution runs against tests, so you see straight away whether your tensor matches the expected output.

• Work through the PyTorch Sheet in order, because the training loop and data loading sections build on the layers and losses before them.
• Rewrite every snippet on this page from memory, write down the shape, dtype or gradient you expect after each line, and then run it to check.
• Write a complete training loop from a blank file at least once a week, including validation, checkpointing and a device switch, until you can do it without stopping.
• Fill the gaps in autograd with Build Micrograd from Scratch, which has 25 problems in seven sections, and the architecture questions with Cracking Deep Learning, which has 35 problems in six sections.
• Take the Neural Net Internals Sprint, which gives you 60 minutes for two neural network problems, and then the ML Engineer Mock Interview to practise under a timer.

What mistakes do candidates make with PyTorch in interviews?

The mistakes interviewers see most often are code that runs but trains the wrong thing, such as layers that never reach the optimizer, gradients that add up across steps, and evaluation that still tracks gradients. Each one produces output without an error, which is why interviewers ask about them directly.

• Storing layers in a plain Python list hides their parameters from the module, so the optimizer never updates them and the checkpoint never saves them, and the fix is nn.ModuleList.
• Calling model.eval() and assuming it turns off gradients wastes memory during evaluation, because only torch.no_grad or torch.inference_mode stops autograd from recording.
• Accumulating gradients over several batches without dividing each loss by the number of steps makes every update several times larger than intended.
• Calling view on a tensor that has just been transposed raises an error, and calling reshape there silently makes a copy, so say which one you want and why.
• Creating a tensor inside forward without device=x.device works on a CPU and fails as soon as the model moves to a GPU.
• Calling scheduler.step() before optimizer.step() shifts the whole schedule by one step and triggers a warning that interviewers notice.
• Comparing tensors with == in a test fails on tiny rounding differences, so use torch.allclose when you check a custom layer against the built-in one.

The how to learn PyTorch guide lists the mistakes beginners make while learning, such as applying softmax before nn.CrossEntropyLoss and keeping loss tensors alive in a logging list.

Practice this

Then work through the full PyTorch Sheet, or take the Neural Net Internals Sprint to answer these questions under a timer.

Related guides

• How to learn PyTorch
• Deep learning interview questions
• NumPy interview questions
• ML coding interview questions
• Which loss function to use
• Training loss not decreasing