torch.compiler.cudagraph_mark_step_begin

torch.compiler.cudagraph_mark_step_begin()

表示即將開始新的推論或訓練迭代。

CUDA Graphs 將釋放先前迭代的張量。 每次調用 torch.compile 都會啟動新的迭代，只要沒有尚未調用的待處理向後傳播即可。

如果該啟發式方法是錯誤的，例如在以下範例中，請使用此 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