comment on torch.compile

fattorib · Nov 11, 2024 · 4feda35 · 4feda35
1 parent f53483a
commit 4feda35
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mamba/fused_cross_entropy.py b/mamba/fused_cross_entropy.py
@@ -73,6 +73,7 @@ def backward(ctx, grad_output):
         return grad_output * weight_grad, grad_output * act_grad, None
 
 
+# torch.compile does a good enough job with the kernel here
 @torch.compile
 def fused_cross_entropy(lm_head_weight, act, labels):
     return CrossEntropyLoopedFused.apply(lm_head_weight, act, labels)