SIGN IN SIGN UP

[autograd] Expose queue_callback as torch.autograd.graph.queue_callback (#193958)

Fixes #66199.

`Variable._execution_engine.queue_callback` is the only way to run code right after the backward pass finishes, but it is a private engine method. We want to use it in [TransformerEngine](https://github.com/NVIDIA/TransformerEngine) to schedule FP8 [delayed scaling](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#Delayed-scaling) amax/scale updates at the end of backward, and we'd rather not depend on an API with no stability guarantees.

This adds `torch.autograd.graph.queue_callback(callback)` as a thin public wrapper over the engine method, with identical semantics — no engine changes. Docs entry, `__all__` and a small test included (+44 lines total).

---
*Drafted with an AI assistant (Claude Code), reviewed by the author.*

Pull Request resolved: https://github.com/pytorch/pytorch/pull/193958
Approved by: https://github.com/soulitzer

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
P
Pawel Gadzinski committed
2bb0ac229e8c3c5781f89d903810f4ed8f609d23
Parent: fd32744
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/24/2026, 5:53:51 PM