fix diloco not defined

PrimeIntellect-ai · Jan 10, 2025 · a22f434 · a22f434
1 parent be3ef20
commit a22f434
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/zeroband/train.py b/src/zeroband/train.py
@@ -169,8 +169,7 @@ def train(config: Config):
         betas=(config.optim.adam_betas1, config.optim.adam_betas2),
     )
 
-    if config.diloco is not None:
-        diloco = Diloco(config.diloco, model, elastic_device_mesh)
+    diloco = Diloco(config.diloco, model, elastic_device_mesh) if config.diloco is not None else None
 
     scheduler = get_scheduler(
         sched_type=config.optim.sched_type,