torch.compiler.allow_in_graph

torch.compiler.allow_in_graph(fn)

自定义编译时将包含在生成图中的函数。 它绕过了对符号化Python代码的所有自省，而是直接将其写入图中。 如果fn是一个可调用对象的列表或元组，它将递归地应用allow_in_graph() 到每个函数，并返回一个包含修改后函数的新列表或元组

Parameters

fn – 一个可调用的对象，表示要包含在图中的函数。

警告

allow_in_graph() 完全跳过装饰函数上的 TorchDynamo，跳过所有 TorchDynamo 安全检查（图中断、处理闭包等）。 因此，使用 allow_in_graph() 时必须非常小心，因为像 AOT Autograd 这样的子系统依赖于 torchdynamo。 如果不小心，这可能会导致正确性问题和非常难以调试的问题。