PyTorch Interview Questions · 2026

PyTorch Interview Questions (2026): Most Asked, With Answers

A candidate interviewing for an ML engineer role at a Series B logistics startup in early 2026 could explain backpropagation on a whiteboard in five minutes flat, chain rule and all. Then the interviewer asked why his training loss had gone up after he added a second call to loss.backward() inside the same step, and he went quiet. He knew the math. He'd never actually noticed that PyTorch accumulates gradients into .grad by default instead of overwriting them, and that optimizer.zero_grad() exists specifically to stop that from happening. PyTorch remains the dominant framework in ML research and is now standard in production teams too, and Stack Overflow's 2025 developer survey tracks it among the most used libraries in the "other frameworks" category across professional developers (Stack Overflow, 2025). Interview loops for ML roles assume you've trained real models with it, not just read the documentation once.

The gap this page is built around: candidates who've only ever called .fit()-style abstractions, or copied training loops from a tutorial without reading them, can usually answer the conceptual questions about tensors and layers just fine. They fall apart the moment a question touches what's actually happening underneath a training loop, why an in-place operation throws a cryptic autograd error, why GPU memory climbs on the second epoch instead of the first, why optimizer.zero_grad() has to run every single step. Those are the questions that tell an interviewer whether you've actually debugged a model at 2am or just run someone else's script successfully once.

This page covers PyTorch interview questions across eight areas: tensors and the autograd engine, building models with nn.Module, training loop mechanics and optimizers, data loading with Dataset and DataLoader, device management and mixed precision, autograd internals and debugging, torch.compile and deployment paths, and distributed training. Code examples are real PyTorch, not pseudocode.

50Questions
Dynamic AutogradCore Concept
PyTorch CodeFormat
Meta AI, 2016Framework Origin

Tensors, autograd, and the computation graph

Every PyTorch loop starts near here, and a shaky answer on how autograd actually tracks operations sets a bad tone for the rest of the interview.

Easy questions

13

A tensor can live on a GPU, a numpy array can't. A tensor can carry a computation graph behind it when requires_grad=True, tracking every operation applied to it so gradients can flow backward later. Numpy has no concept of a gradient at all. Outside of those two things, the day-to-day API is close enough that most numpy code translates almost line for line, indexing, broadcasting, reshaping all work the same way.

python
import torch
import numpy as np

a = np.array([1.0, 2.0, 3.0])
b = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# a has no gradient concept and can't move to GPU
# b tracks operations and can be moved with b.to("cuda")

Setting requires_grad=True tells PyTorch to record every operation performed on that tensor into a computation graph, so that calling .backward() later can walk the graph in reverse and compute gradients with respect to that tensor. It's off by default for plain tensors, on by default for parameters created inside an nn.Module, since those are exactly the values you want gradients for.

python
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x

y.backward()
print(x.grad) # 7.0, since dy/dx = 2x + 3 = 7 at x=2

Because .backward() accumulates into .grad rather than overwriting it. If you skip zero_grad(), the gradients from the previous step are still sitting there, and the new gradients get added on top, so your optimizer step ends up using a sum across however many iterations you forgot to reset. This is exactly the bug the candidate in the intro ran into, and it's one of the most common silent mistakes in a first PyTorch project.

python
for batch in dataloader:
  optimizer.zero_grad()  # reset.grad to None/zero before this step
  output = model(batch.x)
  loss = criterion(output, batch.y)
  loss.backward()     # accumulates into.grad
  optimizer.step()

The accumulation behavior isn't a bug, it's deliberate, and it's the same mechanism gradient accumulation across multiple mini-batches relies on when you skip zero_grad() on purpose for a few steps.

__init__ declares the layers and parameters the module owns. forward defines the actual computation, how input flows through those layers to produce output. Calling the instance directly, model(x), goes through nn.Module.__call__, which runs registered forward hooks and a few bookkeeping steps before invoking forward itself. Calling model.forward(x) directly skips all of that.

python
class Net(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.fc = torch.nn.Linear(10, 1)

  def forward(self, x):
    return self.fc(x)

model = Net()
out = model(x)     # correct, goes through __call__ and any hooks
out = model.forward(x) # works, but silently skips forward hooks

torch.nn.functional.relu is a stateless function, you call it directly on a tensor with no setup. torch.nn.ReLU is a module you instantiate once and call repeatedly, which matters mainly for layers that do carry state, like Dropout or BatchNorm, where the module version needs to exist as an attribute so model.eval() can find and flip it. For a stateless activation like ReLU, either works and it mostly comes down to style, nn.Sequential requires the module form since it only accepts modules in its list.

python
import torch.nn.functional as F

class NetA(torch.nn.Module):
  def forward(self, x):
    return F.relu(self.fc(x)) # functional call, fine for stateless ops

class NetB(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.relu = torch.nn.ReLU() # module form, needed inside nn.Sequential

  def forward(self, x):
    return self.relu(self.fc(x))

Forward pass to get predictions, compute the loss against the targets, zero out old gradients, backward pass to compute new gradients, and step the optimizer to actually update the weights. Get the order of zero_grad and backward wrong and gradients silently accumulate across steps.

python
model.train()
for epoch in range(num_epochs):
  for x, y in dataloader:
    optimizer.zero_grad()
    output = model(x)
    loss = criterion(output, y)
    loss.backward()
    optimizer.step()

state_dict() returns an ordered dictionary mapping every parameter and buffer name to its current tensor value, weights, biases, and things like BatchNorm's running mean, without any of the code that defines the architecture. The standard advice is to save the state dict, not the whole model object with torch.save(model, path), because the latter pickles a reference to the exact class definition and breaks the moment you refactor the file it lived in.

python
torch.save(model.state_dict(), "checkpoint.pt")

# loading requires the architecture to already exist
model = Net()
model.load_state_dict(torch.load("checkpoint.pt", weights_only=True))
model.eval()

weights_only=True matters more than it used to. Recent PyTorch versions warn, and eventually will default to restricting torch.load to plain tensors instead of arbitrary pickled objects, since loading a checkpoint from an untrusted source with full unpickling is a real code-execution risk.

__len__, returning the total number of samples, and __getitem__, returning a single sample given an index. That's the entire contract for a map-style dataset, which is what you'll write nearly all the time. PyTorch's DataLoader calls __getitem__ as many times as it needs per batch and stitches the results together, so any preprocessing you put in __getitem__ runs once per sample, not once per batch.

python
class ImageDataset(torch.utils.data.Dataset):
  def __init__(self, paths, labels, transform):
    self.paths, self.labels, self.transform = paths, labels, transform

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

  def __getitem__(self, idx):
    image = self.transform(load_image(self.paths[idx]))
    return image, self.labels[idx]

.to(device) on an nn.Module moves its parameters and buffers in place, so you don't need to reassign the return value, it mutates the model itself. .to(device) on a tensor is different, it returns a new tensor and doesn't mutate anything, so forgetting to reassign it leaves the original tensor right where it was.

python
model.to(device)     # correct, no reassignment needed, mutates in place

x = x.to(device)     # correct, must reassign
x.to(device)        # bug, does nothing, x is unchanged, still on CPU

Running a model on GPU against a tensor still sitting on CPU throws a clear device-mismatch error, so this bug is usually caught fast. It gets more confusing when it happens three function calls deep inside a data pipeline and the error surfaces somewhere that doesn't obviously point back to the missing reassignment.

.cuda() only ever targets CUDA devices and hardcodes that assumption into the code. .to(device) takes a device string or object, so the same code path works whether device is "cuda", "cpu", "mps" on Apple Silicon, or a specific GPU index, without branching. Almost every modern codebase uses .to(device) with a device resolved once at the top of the script so the rest of the code doesn't need to know or care what hardware it's running on.

python
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

It lets you simulate a larger effective batch size than what actually fits in GPU memory at once, by running several smaller forward and backward passes, accumulating their gradients without stepping the optimizer, and only calling optimizer.step() after however many mini-batches you've decided make up one effective step.

python
accumulation_steps = 4
optimizer.zero_grad()
for i, (x, y) in enumerate(dataloader):
  loss = criterion(model(x), y) / accumulation_steps
  loss.backward()
  if (i + 1) % accumulation_steps == 0:
    optimizer.step()
    optimizer.zero_grad()

torch.save(model.state_dict()) saves only the parameter values as a plain dictionary, portable and stable across refactors as long as the architecture that loads it matches. torch.save(model) pickles the entire object, including a reference to the exact class definition and module path it came from, which means loading it later requires that exact class to still exist, unchanged, at the same import path, or the load fails.

python
torch.save(model.state_dict(), "weights.pt") # recommended, portable
torch.save(model, "full_model.pt")       # works today, brittle across refactors

It seeds PyTorch's CPU and CUDA random number generators, controlling things like weight initialization and dropout masks so they come out the same across runs. It does not, by itself, make GPU computation fully deterministic, because several CUDA operations, some convolution algorithms in particular, use nondeterministic algorithms by default for speed, and can produce very slightly different floating-point results run to run even with the exact same seed.

python
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)

# still not fully deterministic without this, and it costs some speed
torch.use_deterministic_algorithms(True)

Medium questions

27

Every differentiable operation you run on a tensor with requires_grad=True gets recorded as a node in a directed graph, with the output tensor holding a reference back to the operation that produced it through grad_fn. Calling .backward() on a scalar walks that graph from the output backward to every leaf tensor, applying the chain rule at each node, and accumulates the result into each leaf's .grad attribute.

python
x = torch.tensor(3.0, requires_grad=True)
y = x * 2
z = y ** 2

print(z.grad_fn)    # PowBackward
print(z.grad_fn.next_functions) # points back to y's MulBackward

z.backward()
print(x.grad) # d(4x^2)/dx = 8x = 24.0

The graph is built fresh on every forward pass and thrown away after backward unless you explicitly ask to keep it, which is the whole reason PyTorch is called define-by-run rather than define-and-run.

.detach() returns a new tensor that shares the same underlying storage but is severed from the computation graph, so gradients won't flow through it, while everything else around it keeps building the graph normally. torch.no_grad() is a context manager that stops the graph from being built at all for every operation inside it, which is both faster and uses less memory since there's nothing to track in the first place.

python
x = torch.tensor(2.0, requires_grad=True)

y = x.detach() * 3    # y has no grad_fn, graph stops here for y
z = x * 3         # z still tracks x, has a grad_fn

with torch.no_grad():
  w = x * 3       # no graph built at all inside this block

Use .detach() when you need to peel one specific tensor off the graph while keeping the rest live, evaluation loops use no_grad() because nothing in the whole block needs a gradient.

A leaf tensor is one that wasn't created by an operation tracked by autograd, model parameters and any tensor you create directly with requires_grad=True. Intermediate tensors produced by operations on those leaves are not leaves, and PyTorch only stores gradients on leaves by default because that's almost always all you actually need, the values you're about to update with an optimizer step.

python
x = torch.tensor(2.0, requires_grad=True) # leaf
y = x * 3                  # not a leaf, produced by an op

loss = y ** 2
loss.backward()

print(x.grad) # populated, x is a leaf
print(y.grad) # None, and a UserWarning, y is not a leaf

If you genuinely need the gradient on an intermediate tensor, for inspecting activation gradients while debugging, say, call y.retain_grad() before backward, or register a hook on it instead.

Broadcasting lets operations between tensors of different shapes work without an explicit reshape, by stretching dimensions of size 1 to match, following the same rules numpy uses. It's convenient, and it's also a common source of silent shape bugs, because a mismatched shape that happens to be broadcastable doesn't error, it just computes something you didn't intend.

python
predictions = torch.randn(32, 1)  # shape (32, 1), maybe accidentally not squeezed
targets = torch.randn(32)     # shape (32,)

diff = predictions - targets
print(diff.shape) # (32, 32), broadcast into a full pairwise difference matrix
# not the (32,) elementwise difference you probably wanted

This one bites people constantly in loss functions where a stray unsqueezed dimension on either the prediction or the target silently balloons into a much larger tensor, and the loss still runs, it just trains on garbage. Checking .shape right before a loss computation catches this in about two seconds.

Wrapping a tensor in nn.Parameter and assigning it as an attribute of an nn.Module registers it automatically in that module's parameters(), which is what an optimizer iterates over to know what to update, and what state_dict() picks up for saving. A plain tensor attribute, even one with requires_grad=True, is invisible to both, the optimizer never sees it and never updates it.

python
class Layer(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.weight = torch.nn.Parameter(torch.randn(10, 10)) # tracked
    self.scale = torch.randn(10, requires_grad=True)    # not tracked

layer = Layer()
print(list(layer.parameters())) # only weight shows up, scale is invisible

It flips a boolean flag, self.training, on every submodule, and layers that behave differently between training and inference read that flag to decide what to do. Dropout stops zeroing activations and passes everything through unchanged. BatchNorm stops using the current batch's statistics and switches to its running mean and variance accumulated during training instead.

python
model.eval()
with torch.no_grad():
  output = model(x) # dropout is a no-op, batchnorm uses running stats

model.eval() and torch.no_grad() solve two different problems and people mix them up constantly. eval() changes layer behavior. no_grad() stops the graph from being built at all. You almost always want both together at inference time, but forgetting eval() specifically means BatchNorm keeps using batch statistics on a batch of size one during serving, which produces garbage.

nn.ModuleList registers every layer inside it with the parent module, so .parameters(), .to(device), and state_dict() all see through it correctly. A plain Python list of layers stored as an attribute is invisible to all of that, the parent module has no idea those submodules even exist.

python
class Bad(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.layers = [torch.nn.Linear(10, 10) for _ in range(3)] # invisible

class Good(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.layers = torch.nn.ModuleList([torch.nn.Linear(10, 10) for _ in range(3)])

print(len(list(Bad().parameters())))  # 0
print(len(list(Good().parameters()))) # 6, weight+bias for each of 3 layers

The Bad version doesn't error anywhere. It trains, loss goes down a little from whatever layers are still tracked elsewhere, and moving the model to GPU with .to("cuda") silently leaves those three layers on CPU, which throws a device-mismatch error the first time data hits them.

nn.Sequential assumes a strictly linear chain, each layer's output becomes the next layer's only input, with no way to reach back and reuse an earlier tensor. A skip connection, like the identity shortcut in a ResNet block, needs to add an earlier activation back in later, which requires branching logic that a plain sequential stack can't express.

python
class ResidualBlock(torch.nn.Module):
  def __init__(self, dim):
    super().__init__()
    self.fc1 = torch.nn.Linear(dim, dim)
    self.fc2 = torch.nn.Linear(dim, dim)

  def forward(self, x):
    identity = x
    out = torch.relu(self.fc1(x))
    out = self.fc2(out)
    return torch.relu(out + identity) # skip connection, can't express in Sequential

They compute the same value mathematically. nn.CrossEntropyLoss takes raw logits and internally applies log_softmax followed by negative log-likelihood in one fused step. nn.NLLLoss expects you to have already applied log_softmax yourself. The most common bug here is applying softmax or log_softmax to your model's output and then feeding that into CrossEntropyLoss, which silently applies softmax a second time and trains on a wrong, over-smoothed distribution without erroring.

python
logits = model(x) # raw, unnormalized scores

# correct
loss = torch.nn.functional.cross_entropy(logits, targets)

# equivalent, but redundant work and a common source of double-softmax bugs
log_probs = torch.nn.functional.log_softmax(logits, dim=1)
loss = torch.nn.functional.nll_loss(log_probs, targets)

A learning rate scheduler adjusts the optimizer's learning rate over the course of training according to some schedule, cosine annealing smoothly decays it from an initial value down toward zero (or a floor) following a cosine curve, instead of a hard step-drop. scheduler.step() always runs after optimizer.step(), never before, since the scheduler is deciding what the learning rate should be for the next step, not the one that just happened.

python
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

for epoch in range(num_epochs):
  for x, y in dataloader:
    optimizer.zero_grad()
    loss = criterion(model(x), y)
    loss.backward()
    optimizer.step()
  scheduler.step() # after the optimizer, once per epoch here

Gradient clipping rescales the gradient vector if its norm exceeds a threshold, capping how large a single update step can be, without changing its direction. It exists to stop occasional exploding gradients, common in RNNs and in the early steps of training deep transformers, from taking one catastrophic step that destroys days of otherwise stable training.

python
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

You'll see it applied almost unconditionally in transformer training scripts and much more selectively in, say, a small CNN classifier where gradients rarely spike enough for it to matter. Adding it defensively costs almost nothing, so a lot of teams just leave it on everywhere.

With reduction='mean', the loss is averaged over the batch, so the resulting gradient magnitude stays roughly the same no matter how large the batch is, since a bigger batch just gets divided by a bigger number. With reduction='sum', the loss and its gradient scale up directly with batch size, so doubling the batch size effectively doubles your learning rate's real-world effect unless you compensate manually.

python
# gradient accumulation over 4 mini-batches to simulate a 4x larger batch
for i, (x, y) in enumerate(dataloader):
  loss = criterion(model(x), y) / 4  # divide by accumulation steps, matching 'mean' semantics
  loss.backward()
  if (i + 1) % 4 == 0:
    optimizer.step()
    optimizer.zero_grad()

Forgetting to divide by the accumulation count when you switch from a normal loop to gradient accumulation is one of the most common reasons a model that trained fine at batch size 32 suddenly diverges once you accumulate to simulate batch size 128.

num_workers spawns that many separate subprocesses to call __getitem__ in parallel and prefetch batches ahead of when the training loop needs them, so the GPU isn't sitting idle waiting on disk I/O and Python-side preprocessing on the main process. The classic bug is anything non-picklable or stateful captured inside the Dataset, an open file handle, a database connection, a lambda, since each worker is a forked or spawned process and needs to serialize the Dataset object across that boundary.

python
class BadDataset(torch.utils.data.Dataset):
  def __init__(self, path):
    self.file = open(path) # file handle, breaks across worker processes

  def __getitem__(self, idx):
   ...

The fix is almost always to open resources lazily inside __getitem__ or in a per-worker init hook, not in __init__, so each worker process opens its own handle instead of trying to share one that was created before the fork.

Pinned (page-locked) host memory lets the CUDA driver copy data to the GPU asynchronously and faster than it can from regular pageable memory, since the OS can't move pinned pages around mid-transfer. Setting pin_memory=True on the DataLoader allocates batches in that pinned memory, and combined with non_blocking=True on the .to(device) call, the host-to-device copy can overlap with other work instead of blocking the training loop.

python
loader = torch.utils.data.DataLoader(dataset, batch_size=64, pin_memory=True)

for x, y in loader:
  x = x.to("cuda", non_blocking=True)
  y = y.to("cuda", non_blocking=True)

It only matters when you're actually training on GPU and the CPU-to-GPU transfer is a meaningful chunk of your step time. For a small model where the GPU forward and backward pass dominate, the difference is close to invisible.

shuffle=True just randomizes the order samples are drawn in each epoch, uniformly. A custom Sampler gives you control over exactly which indices get drawn and in what order or proportion, which you need the moment your data isn't uniform in a way that matters, class imbalance where you want to oversample the rare class, or distributed training where each process needs a distinct, non-overlapping slice of the dataset.

python
from torch.utils.data import WeightedRandomSampler

# oversample the minority class instead of drawing uniformly
weights = [class_weight[label] for label in dataset.labels]
sampler = WeightedRandomSampler(weights, num_samples=len(dataset), replacement=True)

loader = torch.utils.data.DataLoader(dataset, batch_size=64, sampler=sampler)
# note: sampler and shuffle=True are mutually exclusive, sampler decides the order

Spawning worker processes has real startup overhead, and if a dataset is small enough or preprocessing cheap enough, the main process could have iterated through it faster than the time it takes to fork workers, serialize the dataset object to each of them, and coordinate handing batches back across process boundaries. For a dataset that already fits in memory with lightweight transforms, num_workers=0 running in the main process is often genuinely faster.

This is one of those defaults people copy from a large-scale training script into a small side project without thinking about it, and then wonder why their toy experiment with 2,000 samples runs slower than the same code with num_workers=0.

torch.autocast is a context manager that automatically runs eligible operations in a lower-precision dtype, typically float16 or bfloat16 on GPU, while keeping numerically sensitive operations like reductions and normalization layers in float32, based on an internal list of which ops are safe to downcast. You don't manually cast individual tensors, autocast decides per-operation inside the context.

python
with torch.autocast(device_type="cuda", dtype=torch.float16):
  output = model(x)     # matmuls run in float16
  loss = criterion(output, y) # loss computation stays float32 internally where it matters

The payoff is real, lower-precision matmuls run meaningfully faster on modern GPU tensor cores and use less memory, which is why mixed precision is close to a default in most training recipes now rather than an optional optimization.

It releases GPU memory that PyTorch's caching allocator is holding onto but isn't currently using, back to the CUDA driver, so other processes on the same GPU can use it. It does not free memory that's still referenced by a live Python object, calling it doesn't fix a memory leak, it only returns memory that was already unused but cached for reuse.

python
del model
torch.cuda.empty_cache() # returns freed memory to the driver, doesn't fix leaks

People reach for it as a fix for OOM errors fairly often, and it almost never is one, since if you're actually running out of memory, the problem is that something is still referenced and alive, not that PyTorch is hoarding unused memory unnecessarily. It's genuinely useful when you're sharing a GPU with another process and want to hand back memory your script isn't using anymore between phases.

A static graph, the original TensorFlow 1.x model, gets defined once ahead of time and then executed repeatedly, which lets a compiler optimize the whole graph up front but makes debugging and control flow that depends on runtime values (a different number of RNN steps per batch, say) much more awkward to express. PyTorch builds the graph dynamically, define-by-run, tracing whatever Python code actually executes on each forward pass, which means you can use plain Python if statements and loops and the graph just reflects whatever path was taken that time.

python
def forward(self, x, use_extra_layer):
  x = self.fc1(x)
  if use_extra_layer:    # ordinary Python control flow, graph reflects the branch taken
    x = self.fc2(x)
  return x

The dynamic-versus-static distinction matters less than it used to. torch.compile, added in PyTorch 2.0, captures a graph from that same dynamic Python code and optimizes it ahead of time, getting a lot of the static-graph performance benefit without giving up the flexibility of writing normal Python (PyTorch 2.0 release blog).

Start by actually logging gradient norms per layer instead of guessing, a hook or a loop over model.named_parameters() after .backward() checking param.grad.norm() for each layer tells you immediately whether the problem is concentrated in specific layers (often the earliest ones in a deep stack, for vanishing gradients) or global.

python
loss.backward()
for name, param in model.named_parameters():
  if param.grad is not None:
    print(name, param.grad.norm().item())

From there the fix depends on what you find. Vanishing gradients in a deep network usually point at needing residual connections, a different activation (ReLU family instead of sigmoid/tanh in the middle of a deep stack), or better weight initialization. Exploding gradients usually mean it's time for gradient clipping or a lower learning rate, and if it's happening in the first few steps specifically, a learning rate warmup often fixes it outright.

inference_mode(), added in PyTorch 1.9, does everything no_grad() does and goes further: it also disables version counter tracking on tensors, which is the bookkeeping autograd uses to detect in-place modifications. That extra tracking has real overhead, so tensors created inside inference_mode() run measurably faster than the same code inside no_grad().

python
with torch.inference_mode():
  output = model(x) # faster than no_grad(), but output tensors can't later require grad

The tradeoff: a tensor created inside inference_mode() is permanently marked as an inference tensor and can never be used in an autograd-tracked operation later, even outside the context, which throws an error if you try. For genuine inference-only code paths, model serving, evaluation loops, inference_mode() is the better default. no_grad() is still the right call anywhere you might need to hand a tensor back into a tracked computation afterward.

torch.compile, the headline feature of PyTorch 2.0, traces your model's Python code using a component called TorchDynamo, captures the operations into an optimized graph, and hands that graph to a compiler backend, TorchInductor by default, which generates fused, lower-level kernels (often using Triton) that run faster than PyTorch's normal eager-mode execution (PyTorch 2.0 release blog, 2023).

python
model = Net().to(device)
compiled_model = torch.compile(model)

for x, y in dataloader:
  output = compiled_model(x) # first call triggers compilation, then runs the compiled graph
  loss = criterion(output, y)
  loss.backward()
  optimizer.step()

Unlike TorchScript's tracing, TorchDynamo can fall back to normal eager execution for any piece of Python it can't safely capture, graph breaks, rather than failing outright, which is why torch.compile works on a much wider range of existing model code with little to no rewriting.

TorchScript converts a model into a serialized, self-contained intermediate representation that can be loaded and run entirely outside Python, in a C++ application, on mobile, in an environment where you don't want a Python interpreter in the deployment path at all. torch.compile is squarely a Python-side speed optimization, the compiled model still needs a live Python process to run.

python
scripted = torch.jit.script(model)
scripted.save("model.pt")
# loadable from libtorch in C++, no Python runtime required

If your deployment target genuinely has no Python, embedded, some mobile pipelines, certain C++ inference servers, TorchScript (or increasingly, ONNX export) is still the right tool. If you're serving from a Python process either way, torch.compile gets you speed without the porting headaches TorchScript's tracing and scripting restrictions used to add.

ONNX (Open Neural Network Exchange) is a framework-agnostic model format, exporting a PyTorch model to ONNX lets it run in runtimes that have nothing to do with PyTorch, ONNX Runtime, TensorRT, mobile and edge inference engines, without needing PyTorch installed at all on the serving side.

python
torch.onnx.export(
  model, dummy_input, "model.onnx",
  input_names=["input"], output_names=["output"],
  dynamic_axes={"input": {0: "batch_size"}} # let batch size vary at inference time
)

The most common breakage: export traces the model with one concrete example input, so any control flow that depends on a runtime value, an if branch keyed on a tensor's contents, a variable number of loop iterations, gets baked into the export as whatever path that one example happened to take, silently, with no warning that other inputs would have taken a different path. Custom or unusual ops that don't have a defined ONNX equivalent are the other frequent source of export failures, usually surfacing as an explicit "op not supported" error rather than silent wrong behavior.

nn.DataParallel runs in a single process, splits each batch across GPUs, and gathers outputs back to one GPU (usually GPU 0) to compute the loss and combine gradients, which makes that one GPU a serialization bottleneck and wastes its memory holding gathered outputs from every other GPU. DistributedDataParallel runs one process per GPU, each with its own full copy of the model, and only synchronizes gradients across processes during backward, with no single GPU acting as a bottleneck.

python
# DDP setup, one process per GPU
torch.distributed.init_process_group(backend="nccl")
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

PyTorch's own documentation recommends DDP over DataParallel even for single-machine, multi-GPU training, because DP's single-process, GIL-bound design leaves real performance on the table that DDP's process-per-GPU design doesn't hit.

The backend determines which communication library actually moves tensors between processes during collective operations like all-reduce. nccl is GPU-optimized, built by NVIDIA specifically for fast GPU-to-GPU communication, and is the right choice for essentially all multi-GPU training. gloo is CPU-friendly and works without any GPU-specific hardware support, useful for CPU-only distributed training or for debugging distributed code on a machine without GPUs at all.

python
torch.distributed.init_process_group(backend="nccl") # multi-GPU training
torch.distributed.init_process_group(backend="gloo")  # CPU-only or debugging

.view() returns a new tensor that shares the same underlying storage as the original, requiring the original to be contiguous in memory, and it fails with an error if the requested shape isn't compatible with that storage layout. .reshape() tries to do the same zero-copy view when possible, but falls back to copying the data into a new contiguous layout when it isn't, so it always succeeds where .view() would throw.

python
x = torch.randn(2, 3)
y = x.transpose(0, 1)   # non-contiguous view

y.view(6)         # RuntimeError, not contiguous
y.reshape(6)        # works, copies data internally if needed
y.contiguous().view(6)   # also works, explicit about the copy

.reshape() is the safer default when you're not sure whether the tensor is contiguous. .view() is worth using deliberately when you specifically want to guarantee no hidden copy is happening, and want an explicit error if your assumption about contiguity turns out to be wrong.

Hard questions

10

PyTorch frees the intermediate buffers used to compute gradients as soon as .backward() finishes, to save memory, on the assumption you won't need that graph again. Calling .backward() a second time on the same graph throws RuntimeError: Trying to backward through the graph a second time, because the values it needs to reuse are gone.

python
loss = model(x).sum()
loss.backward()
loss.backward() # RuntimeError, buffers were freed after the first call

# fix: keep the graph alive if you genuinely need to backward twice
loss = model(x).sum()
loss.backward(retain_graph=True)
loss.backward() # works now, but costs extra memory the whole time

The usual fix in production code isn't retain_graph=True, it's realizing you should have run a fresh forward pass to build a new graph instead of reusing the old one. retain_graph shows up legitimately in things like multi-task losses computed from a shared backbone across separate backward calls, but reaching for it as a first instinct usually means the actual bug is somewhere else.

Adam keeps a per-parameter running estimate of both the gradient's mean and its variance, and scales each parameter's update by that variance estimate, effectively giving every parameter its own adaptive learning rate. That adaptivity helps it move quickly through poorly-conditioned loss landscapes early in training, which is why it often converges faster and needs less learning-rate tuning than plain SGD.

The generalization gap shows up because Adam's aggressive per-parameter scaling can converge to sharper minima that fit the training data well but sit in a less flat region of the loss landscape, and flatter minima tend to generalize better. This is exactly why SGD with momentum, sometimes with a warmup and cosine decay schedule, is still the default in a lot of vision and some large language model training recipes even though it needs more careful tuning to get moving at all. It's not a universal rule, AdamW with weight decay tuned properly closes a lot of that gap, but the tendency is real enough that it's worth knowing which failure mode you're trading against.

float16 has a much smaller representable range than float32, so small gradient values that are perfectly meaningful in float32 can underflow to zero once cast down, silently killing the training signal for those parameters. GradScaler works around this by multiplying the loss by a large scale factor before the backward pass, which pushes small gradients up into float16's representable range, then unscales the gradients back down before the optimizer actually applies them.

python
scaler = torch.cuda.amp.GradScaler()

for x, y in dataloader:
  optimizer.zero_grad()
  with torch.autocast(device_type="cuda", dtype=torch.float16):
    output = model(x)
    loss = criterion(output, y)

  scaler.scale(loss).backward()  # scale up before backward
  scaler.step(optimizer)      # unscale, check for inf/nan, then step
  scaler.update()         # adjust the scale factor for next time

It's also dynamically adaptive, if scaling produces an overflow (inf or nan gradients), the scaler skips that optimizer step entirely and lowers the scale factor for next time. bfloat16, which has the same exponent range as float32 at the cost of less mantissa precision, doesn't have this underflow problem the same way, which is part of why bf16 training on newer GPUs often skips GradScaler entirely.

The most common cause is holding onto tensors that still carry a live reference to their computation graph across iterations, logging loss directly into a Python list instead of loss.item(), for instance. The graph for a given step is normally freed after backward, but if something outside the training loop keeps a reference to a tensor that's part of it, the whole graph, and every activation it retained for backward, stays pinned in memory and accumulates step after step until the allocator finally runs out.

python
losses = []
for x, y in dataloader:
  loss = criterion(model(x), y)
  losses.append(loss)    # bug, keeps the whole graph alive for every step
  # losses.append(loss.item()) # fix, detaches to a plain Python float
  loss.backward()
  optimizer.step()

It shows up on epoch two specifically because epoch one usually fits inside whatever memory PyTorch's caching allocator had already reserved from earlier allocations, and the leak only becomes visible once enough steps have piled up that the reserved pool actually gets exhausted. Nothing about epoch boundaries is special, it's purely a matter of when the accumulated leak crosses the available memory.

An in-place operation modifies a tensor's underlying data directly instead of returning a new tensor, and PyTorch marks these with a trailing underscore by convention, add_, relu_, mul_. The danger is that autograd, for some operations, needs to look at the original (pre-modification) value of a tensor to correctly compute the gradient during backward. If you overwrite that tensor's data in place before backward runs, the value autograd needed is gone, replaced by something else.

python
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x * 2
y[0] = 100.0 # in-place modification of y, which is needed for x's gradient

loss = y.sum()
loss.backward() # RuntimeError: a variable needed for gradient computation was modified

Not every in-place op breaks things, PyTorch tracks versions internally and will only complain if the specific modified value was actually needed later. But it's unpredictable enough from the outside that the safe default is to avoid in-place ops on any tensor that's part of an active graph unless you specifically know it's safe.

It means exactly what the previous question describes, at some point during the forward pass, a tensor autograd was planning to reuse during backward got overwritten in place before backward actually ran. The error message includes a version counter mismatch, and PyTorch's anomaly detection mode can point you at the exact operation that clobbered the value if the default message isn't specific enough to find it.

python
with torch.autograd.set_detect_anomaly(True):
  loss.backward()
# raises the error at the forward operation that caused the problem,
# instead of only at the backward call where it's discovered

The two usual culprits in real code, both worth checking first before turning on anomaly detection: reusing a buffer tensor across iterations and modifying it in place instead of allocating a fresh one each time, and calling an in-place activation like relu_ or dropout_ where the module's residual or skip connection still needs the pre-activation value later in the graph. Anomaly detection makes every backward pass noticeably slower, so it's a debugging tool, not something you leave on in a real training run.

A tensor hook registered with .register_hook() fires during the backward pass specifically, with access to the gradient flowing through that exact tensor, letting you inspect or even modify the gradient in place before it continues propagating backward. A module's register_forward_hook fires during the forward pass and only sees inputs and outputs, it has no visibility into gradients at all, since backward hasn't happened yet when it runs.

python
activations = {}

def save_grad(name):
  def hook(grad):
    activations[name] = grad
  return hook

y = model.layer2(model.layer1(x))
y.register_hook(save_grad("layer2_output"))

loss = criterion(y, target)
loss.backward()
print(activations["layer2_output"]) # gradient of the loss w.r.t. y, captured mid-backward

People reach for tensor hooks when debugging vanishing or exploding gradients at a specific point deep inside a model, since printing .grad only works reliably on leaf tensors, and an intermediate activation isn't one unless you explicitly call .retain_grad() on it or attach a hook instead.

Compilation itself has real up-front cost, the first call through a compiled function triggers tracing and codegen, which can take anywhere from seconds to minutes depending on model size. For a model that only runs a handful of forward passes total, in a short-lived script or a small evaluation job, that compilation overhead can dwarf whatever speedup the compiled kernels would have given you over the run's lifetime.

The other common failure mode is data-dependent control flow and constantly changing input shapes, both of which trigger frequent recompilation, since TorchDynamo has to retrace when it hits a graph break or a shape it hasn't specialized for yet. A training loop where batch size or sequence length changes every single step can end up recompiling on nearly every iteration, which is strictly slower than never compiling at all. Setting dynamic=True helps in some of those cases by telling the compiler to generate shape-agnostic code up front, at some cost to the peak speedup for any single fixed shape.

DDP replicates the full model, optimizer state, and gradients on every GPU, which means a model has to actually fit in a single GPU's memory before you can even start. Fully Sharded Data Parallel shards the model's parameters, gradients, and optimizer state across GPUs instead of replicating them, and each GPU only materializes the full parameters for a given layer briefly, right when it's needed for that layer's forward or backward computation, then discards them again.

python
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model) # parameters, gradients, and optimizer state sharded across GPUs

The tradeoff is more communication, since sharded parameters need to be gathered across GPUs on the fly rather than sitting locally the whole time, but for a model too large to replicate fully on one GPU, FSDP is often the difference between training at all and not.

vmap, part of PyTorch's function transform API, vectorizes a function that was written for a single example so it runs across a batch dimension without you rewriting the function itself or falling back to a Python loop over individual examples. Instead of looping in Python and paying per-iteration overhead for every single sample, vmap pushes the batching down to the actual tensor operations, running as one batched kernel call.

python
from torch.func import vmap, grad

def compute_loss(params, x, y):
  pred = params @ x
  return ((pred - y) ** 2).sum()

per_sample_grad = vmap(grad(compute_loss), in_dims=(None, 0, 0))
grads = per_sample_grad(params, batch_x, batch_y) # gradient per sample, no Python loop

The concrete case where this matters most is per-sample gradients, useful for things like differentially private training or certain meta-learning setups, which used to require a genuinely slow per-example Python loop before vmap made it a single vectorized call.

How to prepare for a PyTorch interview in 2026

Skip another pass through the API reference. Write a training loop from scratch, no copy-paste, for a small model on a real dataset, and force yourself to explain every line out loud as you write it: why zero_grad comes before backward, why eval() matters at inference time, what state_dict actually saves. Then break it on purpose. Comment out zero_grad and watch the loss behave strangely. Add an in-place op somewhere and read the actual error message instead of just deleting the line that caused it. The debugging instinct interviewers are testing for only comes from having broken something and fixed it yourself.

Across ML engineering mock interviews run through LastRoundAI, questions about autograd internals, in-place operations, leaf tensors, gradient accumulation, trip up candidates more often than architecture questions do, even candidates who can describe a transformer block from memory without hesitation. Architecture is what people study on purpose from papers and blog posts. Autograd mechanics are what you're expected to have absorbed just from writing enough training code, and plenty of candidates with real project experience never had to look closely enough to actually absorb it.

One more thing worth knowing heading into 2026: torch.compile has matured enough that "does your training loop support compilation" is turning into a real interview question in its own right, not just a curiosity. Being able to explain what a graph break is and why one specific pattern in your code triggers one is a genuinely useful, current signal, and it's the kind of detail that separates someone who trained models on 2.0 from someone whose PyTorch experience is still stuck on 1.x habits, per the framework's official documentation on the underlying mechanics (PyTorch documentation, Autograd mechanics).

Get the reps in before the real thing

Explaining autograd on a whiteboard is not the same as defending a training loop out loud once an interviewer starts asking what happens if you remove one line. LastRoundAI's mock interview mode runs live coding rounds with real-time follow-up questions in your browser, and the free plan includes 15 credits a month that reset monthly rather than piling up unused. Starter is $29/mo (or $19/mo on the 3-month plan) if a handful of sessions isn't enough runway.

Once your answers hold up under a follow-up, the slower part of the job hunt is usually just getting in front of enough ML and applied research roles that actually test PyTorch instead of treating it as a resume keyword. Auto-Apply queues tailored applications for your review, 10 a month on the free plan, up to 400 a month on the Ultimate plan, and nothing goes out until you approve it.

Questions about either product go to contact@lastroundai.com. That's the only inbox we check.

How this list was built

Worth being straight about where these questions come from, because plenty of pages in this category are not. The set was compiled from a research pass across official documentation, vendor release notes, published engineering writing and public discussion of hiring processes, then cross-checked against the current version of each technology so nothing here describes behaviour that has since changed.

What that means in practice: these are the questions the material supports as reasonable and current for this role, not a transcript of any one company's loop. We have not sat in on your interview and we are not going to claim we have. Treat the list as well-sourced preparation rather than a leaked question bank, and expect your panel to phrase things their own way.

If you spot something out of date, tell us at contact@lastroundai.com and we will fix it.

Frequently asked questions

How long does it take to prepare for a PyTorch interview?

If you already work with PyTorch day to day, a focused week on the areas you avoid in practice is usually enough. Coming in cold, expect three to four weeks. The gap is rarely knowledge; it is being able to explain something you normally just use.

What PyTorch topics come up most often?

Interviewers concentrate on the parts that cause production incidents rather than the parts that are pleasant to learn. Expect the fundamentals to be assumed and the follow-up questions to sit one layer below what a tutorial covers.

Do I need hands-on PyTorch experience to pass?

It shows quickly either way. Textbook answers hold up until the interviewer asks what you did when it broke, and that is usually the question that separates candidates. A small real project you can discuss honestly beats a longer list of familiarity claims.

Is PyTorch still worth learning in 2026?

For interview purposes the question is really whether the teams you are targeting use it, which is worth checking against their actual job postings rather than general popularity rankings. Where it is in use it tends to be deeply embedded and slow to replace.

Leave a Reply

Your email address will not be published. Required fields are marked *