From 9f8cb2eb484c0d1ad65ea1c8f1b8e22b81794783 Mon Sep 17 00:00:00 2001
From: Tianlin Liu <tianlinliu@google.com>
Date: Sun, 10 Jul 2022 07:25:02 -0700
Subject: [PATCH] internal

PiperOrigin-RevId: 460055074
---
 vmoe/moe.py | 1 +
 1 file changed, 1 insertion(+)

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: