torch.compiler.cudagraph_mark_step_begin

torch.compiler.cudagraph_mark_step_begin()[源代码]

表示新一轮的推理或训练即将开始。

CUDA 图会释放前一迭代中的张量。每当调用 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