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