torch._dynamo
Warning:
This module is an early prototype and is subject to change.
torch._dynamo.allow_in_graph(fn)
Customize which functions TorchDynamo will include in the generated
   graph. Similar to torch.fx.wrap().
  torch._dynamo.allow_in_graph(my_custom_function)

  @torch._dynamo.optimize(...)
  def fn(a):
      x = torch.add(x, 1)
      x = my_custom_function(x)
      x = torch.add(x, 1)
      return x

  fn(...)

Will capture a single graph containing my_custom_function().
torch._dynamo.disallow_in_graph(fn)
Customize which functions TorchDynamo will exclude in the generated
   graph and force a graph break on.
  torch._dynamo.disallow_in_graph(torch.sub)

  @torch._dynamo.optimize(...)
  def fn(a):
      x = torch.add(x, 1)
      x = torch.sub(x, 1)
      x = torch.add(x, 1)
      return x

  fn(...)

Will break the graph on torch.sub, and give two graphs each with