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