torch.compiler.cudagraph_mark_step_begin

torch.compiler.cudagraph_mark_step_begin()

指示一個新的推理或訓練迭代即將開始。

CUDA Graphs 將會釋放前一個迭代的張量。每次呼叫 `torch.compile` 時都會啟動一個新的迭代，前提是沒有待處理且尚未呼叫的 backward 操作。

如果這個啟發式判斷不正確，例如在下面的示例中，請手動使用此 API 進行標記。

@torch.compile(mode="reduce-overhead")
def rand_foo():
    return torch.rand([4], device="cuda")


for _ in range(5):
    torch.compiler.cudagraph_mark_step_begin()
    rand_foo() + rand_foo()

更多詳情，請參閱 torch.compiler_cudagraph_trees