apply jack comments

Co-authored-by: Jackmin801 <[email protected]>
PrimeIntellect-ai · Dec 6, 2024 · 30a0f23 · 30a0f23
1 parent 835f5f3
commit 30a0f23
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/zeroband/train.py b/src/zeroband/train.py
@@ -19,8 +19,7 @@
 from zeroband.diloco import Diloco, DilocoConfig
 from zeroband.comms import ElasticDeviceMesh
 from zeroband.loss import cross_entropy_max_z_loss
-from zeroband.optimizers import AdamConfig, get_optimizer
-from zeroband.optimizers.muon import MuonConfig
+from zeroband.optimizers import AdamConfig, MuonConfig, get_optimizer
 
 from zeroband.utils import (
     FakeTokenizer,