fix local pg

PrimeIntellect-ai · Nov 21, 2024 · e4b65ce · e4b65ce
1 parent f0245f1
commit e4b65ce
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/zeroband/diloco.py b/src/zeroband/diloco.py
@@ -68,12 +68,11 @@ def __init__(
             # just force compilation
 
         self.pccl_communicator = pccl_communicator
-        self.cpu_local_mesh = init_device_mesh("cpu", mesh_shape=(self.local_pg.size(),))
-
-        self._logger = get_logger()
         self.world_info = get_world_info()
+        self._logger = get_logger()
 
         self._init_offloaded_optimizer(model=model)
+        self.cpu_local_mesh = init_device_mesh("cpu", mesh_shape=(self.world_info.world_size,))
 
     @torch.no_grad()
     def _init_offloaded_optimizer(self, model):