Update cross_entropy.py

megatron/core/tensor_parallel/cross_entropy.py

@@ -120,8 +120,8 @@ class VocabParallelCrossEntropy:
 
 
 class _VocabParallelCrossEntropy(torch.autograd.Function):
-    @torch.compile(mode='max-autotune-no-cudagraphs')
     @staticmethod
+    @torch.compile(mode='max-autotune-no-cudagraphs')
     def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0):
         """Vocab parallel cross entropy forward function."""