diff --git a/lightglue/lightglue.py b/lightglue/lightglue.py
index 4d6b9d7..14e6a61 100644
--- 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