Add annotation

mzusman · Oct 29, 2024 · 11141de · 11141de
1 parent 994e287
commit 11141de
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py
@@ -18,6 +18,7 @@
 
 
 # Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
+@CustomOp.register("mamba_mixer")
 class MambaMixer(CustomOp):
     """
     Compute ∆, A, B, C, and D the state space parameters and compute