torch.compiler.cudagraph_mark_step_begin

torch.compiler.cudagraph_mark_step_begin()[source][source]

表示即将开始新的推理或训练迭代。

CUDA 图将释放先前迭代的张量。每次调用 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。