快捷方式

torch.compiler.cudagraph_mark_step_begin

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

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

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获取问题解答

查看资源