remove num experts

microsoft · Jul 26, 2024 · ee86648 · ee86648
1 parent 378db20
commit ee86648
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mttl/models/containers/selectors/poly_selector.py b/mttl/models/containers/selectors/poly_selector.py
@@ -38,7 +38,10 @@ def __init__(self, **kwargs) -> None:
         self.n_tasks = len(self.config.task_names) if self.config.task_names else 0
 
         # We add an extra task for the default (average) expert if not found
-        shape = (self.n_tasks + 1, self.config.n_splits, self.config.num_experts)
+        shape = (
+            self.n_tasks + 1,
+            self.config.n_splits,
+        )
         self.module_logits = nn.Parameter(torch.empty(*shape).uniform_(-1e-3, 1e-3))
 
         if self.n_tasks == 0: