diff --git a/torchft/local_sgd.py b/torchft/local_sgd.py
index 6956907..b325b85 100644
--- a/torchft/local_sgd.py
+++ b/torchft/local_sgd.py
@@ -16,7 +16,6 @@
 
 import torch
 from torch import nn, optim
-
 from torch.nn.parameter import Parameter
 from torch.optim.optimizer import Optimizer
 
diff --git a/torchft/local_sgd_test.py b/torchft/local_sgd_test.py
index 10db999..7872fc2 100644
--- a/torchft/local_sgd_test.py
+++ b/torchft/local_sgd_test.py
@@ -11,7 +11,7 @@
 import torch
 from torch import nn, optim
 
-from torchft.local_sgd import DiLoCo, DiLoCoOptimizer, LocalSGD
+from torchft.local_sgd import DiLoCo, LocalSGD
 from torchft.manager import Manager
 
 
diff --git a/torchft/manager_integ_test.py b/torchft/manager_integ_test.py
index bb0d1ff..3f5aa6d 100644
--- a/torchft/manager_integ_test.py
+++ b/torchft/manager_integ_test.py
@@ -1,8 +1,8 @@
 import logging
 import threading
 import time
-from concurrent.futures import as_completed, ThreadPoolExecutor
-from contextlib import contextmanager, ExitStack
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field
 from datetime import timedelta
 from typing import Dict, Generator, List, Optional, Protocol, Set, Tuple