Skip to content

Instantly share code, notes, and snippets.

@mlazos
Created March 20, 2024 21:20
Show Gist options
  • Save mlazos/ff6c8ead6bb0f4b17e7d2342317feb14 to your computer and use it in GitHub Desktop.
Save mlazos/ff6c8ead6bb0f4b17e7d2342317feb14 to your computer and use it in GitHub Desktop.
from torch._dynamo.decorators import mark_static_address
@torch.compile(mode="reduce-overhead")
def foo(inp):
return inp + 1
inp = torch.rand([20, 20], device="cuda")
mark_static_address(inp, guard=False)
foo(inp)
inp = torch.ones([20, 20], device="cuda")
foo(inp)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment