Skip to content

Instantly share code, notes, and snippets.

@shashankprasanna
Created April 20, 2023 10:21
Show Gist options
  • Save shashankprasanna/bdd7b553752a9588d07a09d307995e24 to your computer and use it in GitHub Desktop.
Save shashankprasanna/bdd7b553752a9588d07a09d307995e24 to your computer and use it in GitHub Desktop.
def inspect_backend(gm, sample_inputs):
code = gm.print_readable()
with open("forward.svg", "wb") as file:
file.write(FxGraphDrawer(gm,'f').get_dot_graph().create_svg())
return gm.forward
torch._dynamo.reset()
compiled_f = torch.compile(f, backend=inspect_backend)
x = torch.rand(1000, requires_grad=True).to(device)
out = compiled_f(x)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment