diff --git a/reference_algorithms/paper_baselines/momentum/jax/submission.py b/reference_algorithms/paper_baselines/momentum/jax/submission.py
index b2db0c728..dc101896b 100644
--- a/reference_algorithms/paper_baselines/momentum/jax/submission.py
+++ b/reference_algorithms/paper_baselines/momentum/jax/submission.py
@@ -1,7 +1,7 @@
 """Submission file for a SGD with HeavyBall momentum optimizer in Jax."""
 
 import functools
-from typing import Callable, Dict, Iterator, List, Tuple, Any
+from typing import Any, Callable, Dict, Iterator, List, Tuple
 
 from flax import jax_utils
 import jax
diff --git a/reference_algorithms/paper_baselines/momentum/pytorch/submission.py b/reference_algorithms/paper_baselines/momentum/pytorch/submission.py
index 533e9fed4..52aba82bf 100644
--- a/reference_algorithms/paper_baselines/momentum/pytorch/submission.py
+++ b/reference_algorithms/paper_baselines/momentum/pytorch/submission.py
@@ -1,6 +1,6 @@
 """Submission file for a SGD with HeavyBall momentum optimizer in PyTorch."""
 
-from typing import Callable, Dict, Iterator, List, Tuple, Any
+from typing import Any, Callable, Dict, Iterator, List, Tuple
 
 from absl import logging
 import optax
diff --git a/reference_algorithms/paper_baselines/nesterov/jax/submission.py b/reference_algorithms/paper_baselines/nesterov/jax/submission.py
index f79bc34b4..e47c7fa0c 100644
--- a/reference_algorithms/paper_baselines/nesterov/jax/submission.py
+++ b/reference_algorithms/paper_baselines/nesterov/jax/submission.py
@@ -1,7 +1,7 @@
 """Submission file for a SGD with Nesterov momentum optimizer in Jax."""
 
 import functools
-from typing import Callable, Dict, Iterator, List, Tuple, Any
+from typing import Any, Callable, Dict, Iterator, List, Tuple
 
 from flax import jax_utils
 import jax
diff --git a/reference_algorithms/paper_baselines/nesterov/pytorch/submission.py b/reference_algorithms/paper_baselines/nesterov/pytorch/submission.py
index 330e344c1..442949866 100644
--- a/reference_algorithms/paper_baselines/nesterov/pytorch/submission.py
+++ b/reference_algorithms/paper_baselines/nesterov/pytorch/submission.py
@@ -1,6 +1,6 @@
 """Submission file for a SGD with Nesterov momentum optimizer in PyTorch."""
 
-from typing import Callable, Dict, Iterator, List, Tuple, Any
+from typing import Any, Callable, Dict, Iterator, List, Tuple
 
 from absl import logging
 import optax
diff --git a/reference_algorithms/paper_baselines/sam/pytorch/submission.py b/reference_algorithms/paper_baselines/sam/pytorch/submission.py
index 967d53549..15b6b6858 100644
--- a/reference_algorithms/paper_baselines/sam/pytorch/submission.py
+++ b/reference_algorithms/paper_baselines/sam/pytorch/submission.py
@@ -1,6 +1,6 @@
 """Submission file for a SAM optimizer with warmup+cosine LR in PyTorch."""
 
-from typing import Callable, Dict, Iterator, List, Tuple, Any
+from typing import Any, Callable, Dict, Iterator, List, Tuple
 
 from absl import logging
 import torch