fix compilation for torch v2.2.1 (#124)

Mark cudagraph begin in compilation. Fixes #123
cvg · Apr 16, 2024 · edb2b83 · edb2b83
1 parent 075ae42
commit edb2b83
Showing 1 changed file with 1 addition and 0 deletions.
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