From 98365232f2df77f6b1513acb058854f9ed2ef020 Mon Sep 17 00:00:00 2001 From: sami jaghouar Date: Fri, 10 Jan 2025 08:36:41 +0000 Subject: [PATCH] fix diloco not defined --- src/zeroband/train.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/zeroband/train.py b/src/zeroband/train.py index 25cd424a..ad86608e 100644 --- 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,