Never branch on a tensor value you cannot reproduce

Run this over your own model source before you read the rest of this note. It needs none of my code, my model or my results, and whatever it returns is yours.

grep -rn -E "if[[:space:]]+.*\.(any|all|item|sum|mean|max|min)\(\).*:" src/ \
  | grep -v __pycache__

It finds lines where an if is followed by one of seven named calls on a tensor. It will also match comments and strings, so read the hits. Most of what it returns will be harmless. The rest decides the shape of your program on the basis of a floating-point value, and if the number you have to defend was produced by comparing one run against another, that is a difference neither run will show you. I have not caught one of those. The one I caught crashed, which is the lucky version, and it is the one this note is about.

A branch on a tensor value is not really a branch

A Python if is a fork in the program. A comparison against a tensor is a measurement. Put one inside the other and a floating-point value now decides which program runs, and a floating-point value computed on a GPU is not promised to come back identical when you compute it again.

PyTorch publishes the list of operations where that is expected rather than surprising, and what most of them have in common is atomic accumulation: many threads add into the same address, the order they arrive in is not fixed from one launch to the next, and floating-point addition is not associative, so the same inputs can land on a value that differs in its last bits. torch.use_deterministic_algorithms(True) is the switch that makes a listed operation either take a deterministic path or raise rather than run, and it is a switch you throw for a whole run, not for one line. A value you read off a tensor is only as reproducible as the operations upstream of it, and that published list is where you find out which of those are not.

Almost always that is harmless. Inside an if it is not. A quotient that lands at 0.9999999 on the way forward and 1.0000001 when it is computed again clamps to 0.9999999 in one pass and to exactly 1.0 in the other, so a comparison against 1.0 is true in one pass and false in the other. That is not a rounding difference. It is a different program, and nothing in either one knows it took a different route. Nobody dumped the tensor, so read those two figures as the mechanism written out rather than as a measurement.

What that did to a training run

In April 2026, roughly 1,150 steps into a training run of mine, one of these branches took two different routes inside a single step: 162 tensors saved on the way through and 161 on the way back.

torch.utils.checkpoint: A different number of tensors was saved during the original forward and recomputation.
Number of tensors saved during forward: 162
Number of tensors saved during recomputation: 161

Same code, same weights, same batch, same process. The run stopped where it stood and the model state came through intact. That run is not the one this site's scorecard reports on: this was April 2026, the flagship run was lit in July 2026, and the line in my record section about a training loop that did not crash belongs to that later and separate run.

Gradient checkpointing buys memory by discarding the intermediate activations of a wrapped function and running that function again during the backward pass to rebuild them. The arrangement rests on the second run doing what the first run did. The non-reentrant implementation, the one you get with use_reentrant=False, counts the tensors each pass saved and refuses to continue when the counts disagree. Something had created a tensor in one pass and not in the other.

What created it was a norm cap: a module holding the norm of one intermediate tensor down against the norm of another, applying the cap inside a guard, the way anyone writes a guard.

# The shape of the bug, in generic terms. The forward this sits in is wrapped
# in torch.utils.checkpoint.checkpoint(..., use_reentrant=False).

with torch.no_grad():
    ratio = x.detach().float().norm(dim=-1) / y.detach().float().norm(dim=-1).add_(1e-8)
    scale = torch.clamp(cap / ratio, max=1.0)

if (scale < 1.0).any():                      # the whole bug is on this line
    x = x * scale.unsqueeze(-1).to(x.dtype)
    self.cap_fired = True

The branch never entered the autograd graph. scale is computed inside torch.no_grad(), and it is easy to read a no_grad block and conclude that nothing in it can reach the backward pass. What this reached was not a gradient. It was the shape of the recorded computation, which is the thing the non-reentrant implementation counts. Somewhere in that step the comparison came back true on one pass and false on the other, the multiply ran once instead of twice, one extra tensor was saved, and the counter fired.

The timing is the part that transfers. With nothing near the threshold, every scale is 1.0 and the comparison is false in both passes. With everything past it, true in both. What breaks it is the regime in between, some positions marginal and others not, which a training run walks into gradually as its activations move. The branch was there the whole time, as it had to be, since a run does not reload its own source. The conditions for it arrived later, and they arrived visibly: the crash came shortly after the counter tracking how much of the model was exercising the cap stepped up from a few to noticeably more. That is the one thing here anybody could have watched for.

Why this one was loud

It crashed, and the crash was luck. Three things had to be true at once, and none of them is a property of the bug.

The branch had to sit inside a checkpointed region, which is close to the only thing in an ordinary training step that runs the same forward twice. It had to decide whether a tensor got created, because tensors are what the comparison is over. And the checkpoint had to be the non-reentrant one, because that is the only implementation that compares anything. Take any of the three away and the branch still flips with nothing standing there to notice.

That distinction is not academic: the fix below keeps a tensor read inside the same checkpointed forward and is safe, because what the read decides is a Python attribute rather than the graph.

What my own record got wrong

The lesson file I wrote at the time says use_reentrant=False did not save me, and offers replayed random numbers as the thing that flag buys. Both halves are wrong, and I would rather correct my own notes in public than write a post about instruments lying while quietly carrying an error out of one.

The saved-tensor count lives in the non-reentrant implementation only. I read torch/utils/checkpoint.py in torch 2.6.0 rather than trust my own note: the comparison is a method on the frame the non-reentrant path builds, and the reentrant class contains no count comparison at all. Under use_reentrant=True the same flip raises nothing and the gradients are quietly built from a forward that ran differently. So the flag did not fail to save me. It is the reason I found out.

Replaying the random number generator is not what that flag buys either. Checkpointing stashes and restores the generator state around the recomputation, so dropout replays identically, and it does that in both implementations under preserve_rng_state, which defaults to on. What none of it can replay is nondeterminism that did not come from the generator.

The version with nothing counting

I am reasoning now rather than reporting, and I would rather say so than let the two blur. What I have is one incident, and in it the bug crashed the run. I am not claiming a case where it went quietly, because I do not have one. The audit below turned up one live site in my model source and it is the one that crashed. I did not run the same pattern over my evaluation code, which is where the quiet version would live, so I am not claiming that is clean either.

The mechanism has no opinion about where you put it, though. A comparison that floating-point noise can flip flips the same way in a scoring threshold, an early-exit test, a filter deciding whether a sample counts, and nothing in those places is counting saved tensors. No exception, no log line, just a number that came out a little different. That is the version that reaches a decision. When you set one result beside another you are assuming that the only thing differing between the two measurements is what you changed, and a data-dependent branch is not held still: not across runs, not across steps, not inside one step. A number produced that way is not wrong in a way you can find by staring harder at it. It is wrong in a way you reach only by asking whether the code that made it does the same thing twice.

The fix

Remove the branch, not the operation. The multiply was already doing nothing wherever the cap was not biting, because the scale is clamped so an untouched position gets multiplied by 1.0. Deleting the conditional changes no result and makes the graph come out the same shape on both passes whatever the values do. It buys that with a multiply that now always runs, which on an already clamped scale is the identity, and that is a trade worth making.

with torch.no_grad():
    ratio = x.detach().float().norm(dim=-1) / y.detach().float().norm(dim=-1).add_(1e-8)
    scale = torch.clamp(cap / ratio, max=1.0)
    self.cap_fired = bool((scale < 1.0).any().item())    # telemetry only

# Unconditional. scale is 1.0 wherever the cap is not biting, so this is the
# identity there, and the graph comes out the same shape on both passes.
x = x * scale.unsqueeze(-1).to(x.dtype)

The tensor-valued comparison survives as a telemetry flag, and it still calls .item(), which is the thing everyone is told not to do in a forward pass. It is safe in that position for a specific reason rather than a general one: it is computed inside the no_grad block and written to a plain Python attribute, so it is not part of the autograd graph. The flag itself can still land differently on the two passes, and the recomputation writes last, which is fine, because nothing downstream of it is a tensor. The rule is not that reading a tensor value is forbidden. The rule is that what you read must not decide the shape of the program.

Triaging what the grep returned

Work the hits in this order.

1. Read the enclosing function, not the line. If it is ever wrapped in a checkpoint, directly or through a parent module, the branch is deciding the shape of the saved-tensor set on two separate executions and it has to go.
2. Keep the hits whose condition is settled before any tensor is read. if self.training: and if self.cap > 0: are decided by Python state, they are identical on both passes, and they were never part of this problem.
3. Ask what the branch decides. Does it decide whether a tensor gets created, or only what gets logged? Creating a tensor changes the graph. Writing a float to a Python attribute does not.
4. Replace what is left with an unconditional operation that is the identity where the condition would have been false. That is what the fix above does, and it is usually available.
5. Do not skip a hit because the branch is dormant today. Dormant is a configuration, and configurations get changed by whoever runs the next experiment.

Then take the pattern for what it is, a floor rather than a sweep. It catches seven named tensor calls written with empty parentheses inside an if condition, and walks straight past the rest:

if t.norm() > 0: ...           # a tensor read it does not name
if bool(t): ...                # no method call at all
if torch.any(t): ...           # a function, not a method
hit = (scale < cap).any()      # read on one line,
if hit: ...                    # branched on the next

A clean result means the cheapest version of this is absent from your code. It does not mean the class is.

The grep is for code you can read. If what you cannot defend is a number rather than a module, there is a blunter check. Run the same evaluation twice, on the same weights and the same inputs, with torch.use_deterministic_algorithms(True) set, and confirm it agrees with itself exactly. Without that switch it can disagree for reasons that have nothing to do with your branch, and you learn nothing. Set CUBLAS_WORKSPACE_CONFIG as well, because on CUDA torch.mm, torch.mv and torch.bmm raise under that switch without it. Disagreement will not tell you where the branch is. It tells you there is something to find, which is more than the number by itself will ever tell you.

On my own model source the audit returned two hits worth reading. One was the branch that had just crashed, live in a checkpointed path. The other was the same pattern in a different module, inside a code path that was switched off by default, so it was not executing. Dormant is not fixed. It is a landmine with the pin still in, and what I wrote down at the time is that if that option is ever turned back on, that site gets rewritten first. That review is dated the same day as the crash, in April 2026, and I have not re-run the pattern since, so read it as of then rather than as of today.

Never let the value of a tensor decide what your program does, unless you have made the arithmetic behind that value reproducible and can say why.

What the crash was worth

More than it cost. It was run-killing and non-destructive: the model state was safe, and what the branch took was the run rather than the weights. It did not take a metric, and it did not take a decision made on the strength of one.

The habit underneath it is one of the four habits the work runs on, the one that says audit the instrument before believing the number. Here a counter inside gradient checkpointing did that for me, on one narrow question, because the counter happened to be there. Nothing enforces it across the rest of a forward pass unless somebody goes and looks.

The Evaluation Teardown is what that looks like when none of it is left to accident, pointed at your harness instead of mine. The instrument gets attacked before the result, and you keep one written finding per metric that names the confound, gives the exact steps to reproduce it, and states what the number becomes once the confound is gone. If nothing is found, you keep that in writing too. This one announced itself with a stack trace. The ones worth paying somebody to find are the ones that will not.