Use with AI · Chapter Application

Audit PyTorch Autograd and Gradient Flow

Trace the computation graph in a real PyTorch training path and find detach, requires_grad, backward, in-place and NaN-gradient failures using repository evidence.

Chapter Application PyTorch: Zero to Hero PyTorch training and autograd code Intermediate
Read the source chapter → View the learning path →

How to use this

  1. Open a repository-aware AI assistant.
  2. Give it access to the repository or files you want reviewed.
  3. Copy the prompt below and run it unchanged first.
  4. Use the evidence it finds to decide what to inspect or change next.
PromptCopy and run against your own project
You are auditing a PyTorch repository for autograd and gradient-flow defects.

Do not start by changing the optimizer, learning rate, architecture or loss function. First establish whether the intended parameters are actually connected to the loss through a valid computation graph.

Trace one complete training step from model input to optimizer.step().

Work in this order:

1. Identify:
   - the loss tensor
   - trainable parameters
   - optimizer parameter groups
   - zero_grad / backward / step ordering
   - gradient accumulation behavior

2. Trace the graph-breaking risks:
   - detach()
   - torch.no_grad()
   - .item() used too early
   - tensor reconstruction via torch.tensor(existing_tensor)
   - conversion through NumPy/Python values
   - requires_grad changes
   - frozen parameters that should be trainable
   - parameters omitted from the optimizer
   - in-place operations that invalidate gradient history

3. Inspect backward semantics:
   - repeated backward through freed graphs
   - retain_graph used to mask design problems
   - gradient accumulation without intentional scaling
   - zeroing gradients at the wrong time
   - multiple losses with unintended graph interaction

4. Inspect numerical gradient failures:
   - NaN / Inf loss
   - NaN / Inf gradients
   - exploding norms
   - mixed precision misuse
   - unsafe divisions/logs/exponentials
   - masking that creates invalid values

5. Classify every finding as:
   - confirmed graph break
   - confirmed optimizer/gradient lifecycle defect
   - confirmed numerical instability
   - suspicious but unproven
   - valid behavior

For each real issue provide:
- file/function
- the exact graph or lifecycle break
- affected parameters
- evidence from the code
- smallest corrective change
- a test or runtime probe proving gradients now flow correctly

Output:

## Training-step graph
Describe the actual loss → backward → parameter update path.

## Findings
Ranked by severity and confidence.

## Minimal corrections
Do not broaden beyond demonstrated failures.

## Verification
Include checks such as parameter grad presence, finite-gradient assertions, gradient norms and before/after parameter deltas.

If autograd is correctly wired, say so and identify which gradient assumptions deserve tests.