fix(env_check): use the amp_bf16 precision for the CPU autocast branch

[Anai-Guo]

Problem

check_accuracy() in torchbenchmark/util/env_check.py builds the list of autocast contexts the accuracy run executes under. The CPU branch is unreachable:

elif (
    tbmodel.dargs.precision == "amp"
    and tbmodel.dargs.precision == "bf16"
    and tbmodel.device == "cpu"
):
    contexts.append(torch.cpu.amp.autocast)

tbmodel.dargs.precision cannot be both "amp" and "bf16" at the same time, so this condition is always False and torch.cpu.amp.autocast is never appended. contexts is what forward_pass / forward_and_backward_pass open via nested(*contexts), so a CPU bf16-autocast accuracy run silently executes in plain fp32 — it reports on a configuration it never actually exercised.

Why amp_bf16 is the intended value

"amp" and "bf16" are two separate entries in AVAILABLE_PRECISIONS, and the mode this branch wants has its own third entry, "amp_bf16". The repo already pairs that value with the CPU device for this exact autocast context, in torchbenchmark/util/extra_args.py:

elif dargs.precision == "amp_bf16":
    assert model.device == "cpu", "amp_bf16 is only supported on cpu device."
    if model.test == "eval":
        model.add_context(lambda: torch.cpu.amp.autocast(dtype=torch.bfloat16))

and is_staged_train_test support is gated the same way at extra_args.py:54 (precision == "amp_bf16" + device == "cpu"). So this is the env_check.py copy of an existing, correct condition.

Fix

elif tbmodel.dargs.precision == "amp_bf16" and tbmodel.device == "cpu":
    contexts.append(torch.cpu.amp.autocast)

Verification

I sliced the contexts = [] … contexts.append(torch.cpu.amp.autocast) block straight out of the file (before and after the patch, via git show HEAD:… vs the working copy — not hand-transcribed) and replayed it over every cell of AVAILABLE_PRECISIONS × {cpu, cuda} × {eval, train}, with torch stubbed so the appended context is identifiable:

total cells: 40   changed: 2   identical: 38
   amp_bf16  cpu  eval     [] -> ['CPU_AUTOCAST']
   amp_bf16  cpu  train    [] -> ['CPU_AUTOCAST']

CPU autocast reachable BEFORE: []
CPU autocast reachable AFTER : [('amp_bf16','cpu',eval), ('amp_bf16','cpu',train)]

The CUDA branch above it is untouched, and 38 of the 40 cells are bit-identical — the only behaviour change is that the two cells extra_args.py already documents now actually get the autocast context.

black (repo pyproject.toml, line-length 88) reports the file unchanged.

🤖 Generated with Claude Code