src/aiia/model/Model.py

@@ -179,7 +179,7 @@ class AIIAExpert(AIIA):
 
 class AIIAmoe(AIIA):
     def __init__(self, config: AIIAConfig, num_experts: int = 3, base_class=AIIABase, **kwargs):
-        super().__init__()
+        super().__init__(config=config, **kwargs)
         self.config = config
 
         # Update the config to include the number of experts.