git pull

hivemind/moe/client/balanced_expert.py

@@ -34,12 +34,12 @@ class BalancedRemoteExpert(nn.Module):
     ):
         super().__init__()
         if uid_prefix.endswith(".0."):
-            logger.warning(f"BalancedRemoteExperts will look for experts under prefix {self.uid_prefix}.0.")
+            logger.warning(f"BalancedRemoteExperts will look for experts under prefix {self.uid_prefix}0.")
         assert len(grid_size) == 2 and grid_size[0] == 1, "only 1xN grids are supported"
         self.dht, self.uid_prefix, self.grid_size = dht, uid_prefix, grid_size
         self.forward_timeout, self.backward_timeout = forward_timeout, backward_timeout
         self.backward_task_size_multiplier = backward_task_size_multiplier
-        self.expert_balancer = ExpertBalancer(dht, key=f"{self.uid_prefix}.0.", update_period=update_period, **kwargs)
+        self.expert_balancer = ExpertBalancer(dht, key=f"{self.uid_prefix}0.", update_period=update_period, **kwargs)
         self._expert_info = None  # expert['info'] from one of experts in the grid
 
     def forward(self, *args: torch.Tensor, **kwargs: torch.Tensor):