fix compilation for torch v2.2.1 (#124)

Phil26AT · web-flow · commit edb2b838efb2 · 2024-04-16T15:33:36.000+02:00
Mark cudagraph begin in compilation. Fixes #123
diff --git a/lightglue/lightglue.py b/lightglue/lightglue.py
@@ -438,6 +438,7 @@ def compile(
                 stacklevel=2,
             )
 
+        torch._inductor.cudagraph_mark_step_begin()
         for i in range(self.conf.n_layers):
             self.transformers[i].masked_forward = torch.compile(
                 self.transformers[i].masked_forward, mode=mode, fullgraph=True

Original file line number	Diff line number	Diff line change
`@@ -438,6 +438,7 @@ def compile(`
`438`	`438`	`stacklevel=2,`
`439`	`439`	`)`
`440`	`440`
	`441`	`+ torch._inductor.cudagraph_mark_step_begin()`
`441`	`442`	`for i in range(self.conf.n_layers):`
`442`	`443`	`self.transformers[i].masked_forward = torch.compile(`
`443`	`444`	`self.transformers[i].masked_forward, mode=mode, fullgraph=True`