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