black-isort

hivemind/optim/__init__.py

@@ -1,6 +1,7 @@
 from hivemind.optim.adaptive import CollaborativeAdaptiveOptimizer
 from hivemind.optim.base import DecentralizedOptimizerBase
 from hivemind.optim.collaborative import CollaborativeOptimizer
-from hivemind.optim.grad_scaler import HivemindGradScaler
+from hivemind.optim.experimental.optimizer import Optimizer
+from hivemind.optim.grad_scaler import GradScaler, HivemindGradScaler
 from hivemind.optim.simple import DecentralizedAdam, DecentralizedOptimizer, DecentralizedSGD
 from hivemind.optim.training_averager import TrainingAverager

hivemind/optim/experimental/optimizer.py

@@ -243,7 +243,6 @@ class Optimizer(torch.optim.Optimizer):
                 f"BEFORE: {self.grad_averager.local_samples_accumulated}, {repr([grad.norm() / self.grad_averager.local_times_accumulated for grad in self.grad_averager._grad_accumulators()])}"
             )
 
-
             need_averaging = self.tracker.global_progress.num_peers > 1
             if need_averaging:
                 try: