25 PyTorch interview questions

Listen for: A strong answer narrows the problem instead of guessing: inspect inputs and labels, check tensor shapes and ranges, reproduce with a small batch, enable torch.autograd.set_detect_anomaly only when needed, check learning rate and initialization, log gradient norms, and isolate the exact operation producing Inf or NaN.

Red flag: They immediately blame PyTorch or the GPU without proposing a reproducible debugging path.

Listen for: Listen for separation of indexing from transformation, lazy loading in __getitem__, deterministic train/validation splits, appropriate num_workers and pin_memory, avoiding GPU work inside Dataset, and awareness that heavy augmentations or slow storage can starve the GPU. They should mention verifying samples visually and testing one batch before training.

Red flag: They load the entire dataset into GPU memory inside the Dataset.

Listen for: They should discuss model.train() versus model.eval(), dropout and batch normalization behavior, torch.no_grad() or torch.inference_mode(), device and dtype consistency, preprocessing differences, random seeds, and whether validation accidentally updates metrics or model buffers. A good answer also checks that the same checkpoint weights are actually loaded.

Red flag: They think eval mode disables gradient computation by itself.

Listen for: Good signals include creating a tiny synthetic batch, asserting input and output shapes, running a forward pass, checking loss is scalar, calling backward, confirming key parameters receive non-None gradients, running one optimizer step, and optionally overfitting a very small batch to prove the loop can learn.

Red flag: They rely only on reading the code and skip runtime checks.

Listen for: They should trace tensor creation sites, move batches and models explicitly to a device, avoid hard-coded cuda calls, create new tensors with x.new_*, torch.zeros_like, or device=existing.device, and test on CPU and GPU. Strong candidates also mention buffers registered with register_buffer for non-parameter tensors.

Red flag: They scatter .cuda() calls throughout the code until the error disappears.

Listen for: Expect model state_dict, optimizer state_dict, scheduler state, epoch or step, best metric, configuration, label mapping, random seed or RNG states when reproducibility matters, and PyTorch version or git commit. They should prefer state_dict over pickling whole objects unless there is a clear reason.

Red flag: They save only the model weights and cannot resume training correctly.

Listen for: A strong answer places autocast around forward and loss computation, uses GradScaler when training with CUDA AMP where appropriate, divides loss or otherwise accounts for accumulation, unscales before clipping, steps and updates scaler at the right interval, handles remainder batches, and logs effective batch size and skipped steps.

Red flag: They clip scaled gradients or step the optimizer on every microbatch by accident.

Listen for: Look for a structured approach: separate input pipeline from model compute, use torch.profiler with schedules and NVTX ranges, inspect CPU time, CUDA kernels, synchronization points, DataLoader wait time, and memory copies. They should test num_workers, pin_memory, prefetching, batch size, transforms, and accidental .item() or print calls.

Red flag: They increase model size or batch size without measuring the bottleneck.

Listen for: They should know torch.compile can help stable tensor programs but may struggle with graph breaks, dynamic control flow, custom ops, unsupported Python side effects, or dynamic shapes. Strong candidates benchmark eager versus compiled modes, validate numerics, cache warmup effects, fallback behavior, and isolate compile settings behind configuration.

Red flag: They treat torch.compile as a guaranteed speed switch for every model.

Listen for: Listen for versioned checkpoint metadata, explicit config schemas, state_dict key migration, strict=False only with logged missing and unexpected keys, tests that load old checkpoints, separating training and inference artifacts, and saving tokenizer or preprocessing state. They should avoid silent partial loads that change model behavior unnoticed.

Red flag: They use strict=False to hide all loading problems.

Listen for: Good answers cover data quality and leakage, baseline frozen-head training, learning-rate choices for backbone versus head, augmentation or regularization, class imbalance, evaluation protocol, checkpoint selection, and whether parameter-efficient methods such as LoRA are appropriate. They should discuss overfitting detection before reaching for a larger model.

Red flag: They immediately unfreeze everything and train with the same settings as pretraining.

Listen for: They should mention mismatched collective calls, uneven batches, exceptions on one rank, DataLoader worker failures, conditional branches that skip parameters, sampler misconfiguration, NCCL environment logs, timeouts, and rank-specific logging. Strong answers include reproducing on two GPUs, checking DistributedSampler.set_epoch, and making all ranks reach the same barriers.

Red flag: They debug only rank zero logs and ignore the other processes.

Listen for: Expect a clear artifact boundary: model weights, architecture or exported graph where appropriate, preprocessing and postprocessing, config, input/output schema, device and dtype expectations, tests with golden examples, and simple versioning. They should compare eager PyTorch, torch.export, TorchScript legacy use, ONNX, or TensorRT based on the consumer.

Red flag: They hand over a notebook and assume the other team can run it unchanged.

Listen for: Strong candidates connect training data, evaluation slices, preprocessing parity, checkpoint lineage, calibration, serving dtype, batching, latency fallbacks, and feature drift. They should propose replay tests, shadow evaluation, prediction logging with privacy constraints, and rollback. The best answers distinguish model bugs from measurement or product integration issues.

Red flag: They focus only on retraining until the offline metric improves again.

Listen for: Look for a small set of stable abstractions: config, model factory, data module, trainer loop, callbacks, logging, checkpointing, and distributed strategy. Strong answers avoid excessive inheritance, keep escape hatches for custom steps, support torch.compile and AMP behind options, and enforce tests around reproducibility and checkpoint compatibility.

Red flag: They propose a large framework that hides all PyTorch details and cannot handle unusual models.

Listen for: They should compare batch size reduction, activation checkpointing, mixed precision or bfloat16, gradient accumulation, optimizer state sharding with FSDP or ZeRO-style tools, parameter-efficient fine-tuning, offloading, sequence length changes, and memory-efficient attention. Good answers discuss throughput, convergence, implementation risk, and debugging cost rather than listing tricks.

Red flag: They enable every memory optimization at once without measuring impact.

Listen for: Strong answers define acceptable tolerances, compare against a full-precision baseline, test sensitive operations such as softmax, normalization, reductions, and loss computation, monitor gradients, handle scaler behavior for fp16, and keep some layers or accumulations in fp32 when needed. They also separate training precision from inference quantization decisions.

Red flag: They assume lower precision is correct if the code runs without errors.

Listen for: They should ask about runtime target, dynamic shapes, custom ops, control flow, preprocessing, postprocessing, required operators, latency constraints, debuggability, and parity tests. Strong candidates know PyTorch 2.x export paths are useful but require validation, and that eager PyTorch may still be the right serving choice.

Red flag: They promise any PyTorch model can be exported cleanly after training.

Listen for: Listen for standards beyond style: deterministic smoke tests, shape assertions, device and dtype handling, no hidden global state, safe checkpointing, clear config ownership, profiler evidence for performance changes, distributed correctness, and separation of experiment code from production paths. Good answers include mentoring and review templates, not gatekeeping.

Red flag: They review only model architecture and ignore data, training, and operational risks.

Listen for: Strong answers include request latency by stage, batch size, queue time, GPU utilization, memory, OOMs, model version, input validation failures, numerical anomalies, fallback rates, and output distribution drift. They plan for warmup, rolling deploys, canaries, health checks, graceful CPU or smaller-model fallback where appropriate, and rollback.

Red flag: They monitor only process uptime and average latency.

Listen for: A strong design covers job submission, environment locking with containers, dataset access, experiment config, logging, checkpoint storage, resume behavior, quotas, and GPU scheduling. It should support DDP/FSDP jobs, PyTorch 2.x options like compile and AMP, and clear ownership boundaries between platform code and model code.

Red flag: They design only a single training script and ignore scheduling, artifacts, and reproducibility.

Listen for: Look for request schema validation, preprocessing placement, batching strategy, model warmup, memory planning, CUDA stream awareness at a high level, dtype choice, versioned artifacts, canary deploys, and observability. Strong answers compare eager PyTorch, torch.compile, exported graphs, and TensorRT/ONNX Runtime without assuming one is always best.

Red flag: They batch everything aggressively without considering tail latency or request deadlines.

Listen for: They should discuss data filtering, prompt or image preprocessing, tokenizer and processor versioning, LoRA or other parameter-efficient tuning, FSDP or DeepSpeed trade-offs, gradient checkpointing, mixed precision, evaluation slices, safety or quality review, and checkpoint merging. Strong designs include failure recovery and cost-aware experiment pruning.

Red flag: They start with full fine-tuning and no baseline or memory plan.

Listen for: Good answers cover sharded input reading, deterministic IDs, batched GPU inference, preprocessing parity, no_grad or inference_mode, checkpoint versioning, partial retries, output format, vector index handoff, and data quality checks. They should plan for throughput without losing traceability from each embedding back to source data.

Red flag: They produce embeddings without recording the model, preprocessing, or source version.

Listen for: Strong candidates propose extracting data, model, training, evaluation, and config modules; adding smoke tests and golden-batch tests; standardizing logging and checkpoints; pinning environments; adding CI for CPU and optional GPU; and documenting experiment reproduction. They keep the first production version small rather than rewriting everything.

Red flag: They either keep the notebook as production code or propose a long rewrite before any value ships.