diff --git a/vmoe/moe.py b/vmoe/moe.py
index c79826f..18b4aba 100644
--- a/vmoe/moe.py
+++ b/vmoe/moe.py
@@ -65,6 +65,7 @@ class BaseDispatcher(abc.ABC):
   combine might not be the same as the ones at dispatch (e.g. if the expert
   changes the shape of the data).
   """
+  combine_weights: Array
 
   @abc.abstractmethod
   def dispatch(self, data: Array) -> Array: