Update jit.py

megatron/core/jit.py

@@ -7,7 +7,7 @@ from megatron.core.utils import is_torch_min_version
 jit_fuser = torch.jit.script
 # nvFuser is deprecated in PyTorch JIT starting from 2.2
 if is_torch_min_version("2.2.0a0"):
-    jit_fuser = torch.compile
+    jit_fuser = torch.compile(mode='max-autotune-no-cudagraphs')
 
 # Decorator to disable Torch Dynamo
 # See: https://github.com/NVIDIA/TransformerEngine/issues/308