diff --git a/apex/__init__.py b/apex/__init__.py
index 5c991871d..406a9a363 100644
--- a/apex/__init__.py
+++ b/apex/__init__.py
@@ -11,11 +11,9 @@
 # so they expect those backends to be available, but for some reason they actually aren't
 # available (for example because they built improperly in a way that isn't revealed until
 # load time) the error message is timely and visible.
-from . import optimizers
-from . import normalization
+from . import normalization, optimizers
 
-
-__all__ = ["optimizers", "normalization"]
+__all__ = ["normalization", "optimizers"]
 
 
 def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int) -> bool:
diff --git a/apex/_autocast_utils.py b/apex/_autocast_utils.py
index 3a92a83f3..a67771812 100644
--- a/apex/_autocast_utils.py
+++ b/apex/_autocast_utils.py
@@ -1,8 +1,8 @@
-from typing import Optional, Sequence
+from collections.abc import Sequence
+from typing import Optional
 
 import torch
 
-
 __all__ = ["_cast_if_autocast_enabled"]
 
 
@@ -12,7 +12,7 @@ def _get_autocast_dtypes() -> Sequence[torch.dtype]:
     return [torch.half]
 
 
-def _get_current_dtype(dtype: Optional[torch.dtype] = None) -> torch.dtype:
+def _get_current_dtype(dtype: torch.dtype | None = None) -> torch.dtype:
     if not torch.is_autocast_enabled():
         return torch.float or dtype
     else:
diff --git a/apex/contrib/bottleneck/__init__.py b/apex/contrib/bottleneck/__init__.py
index 5263e6874..46ae84ab4 100644
--- a/apex/contrib/bottleneck/__init__.py
+++ b/apex/contrib/bottleneck/__init__.py
@@ -1,7 +1,7 @@
 from .bottleneck import Bottleneck, SpatialBottleneck
 from .halo_exchangers import (
-    HaloExchangerNoComm,
     HaloExchangerAllGather,
-    HaloExchangerSendRecv,
+    HaloExchangerNoComm,
     HaloExchangerPeer,
+    HaloExchangerSendRecv,
 )
diff --git a/apex/contrib/bottleneck/bottleneck.py b/apex/contrib/bottleneck/bottleneck.py
index e24a9251b..bc2ed2558 100644
--- a/apex/contrib/bottleneck/bottleneck.py
+++ b/apex/contrib/bottleneck/bottleneck.py
@@ -1,12 +1,11 @@
 import functools as func
 
+import fast_bottleneck
+import nccl_p2p_cuda as inc
 import torch
 from torch import nn
 
 from apex import check_cudnn_version_and_warn
-import fast_bottleneck
-import nccl_p2p_cuda as inc
-
 
 assert check_cudnn_version_and_warn(__name__, 8400)
 
@@ -35,7 +34,7 @@ class FrozenBatchNorm2d(torch.jit.ScriptModule):
     """
 
     def __init__(self, n):
-        super(FrozenBatchNorm2d, self).__init__()
+        super().__init__()
         self.register_buffer("weight", torch.ones(n))
         self.register_buffer("bias", torch.zeros(n))
         self.register_buffer("running_mean", torch.zeros(n))
@@ -171,7 +170,7 @@ def __init__(
         use_cudnn=False,
         explicit_nhwc=False,
     ):
-        super(Bottleneck, self).__init__()
+        super().__init__()
         if groups != 1:
             raise RuntimeError("Only support groups == 1")
         if dilation != 1:
@@ -224,8 +223,6 @@ def __init__(
                 with torch.no_grad():
                     p.data = p.data.permute(0, 2, 3, 1).contiguous()
 
-        return
-
     # Returns single callable that recomputes scale and bias for all frozen batch-norms.
     # This method must be called before cuda graphing.
     # The callable it returns can be called anytime.
@@ -851,7 +848,7 @@ def __init__(
         explicit_nhwc=False,
         spatial_parallel_args=None,
     ):
-        super(SpatialBottleneck, self).__init__()
+        super().__init__()
         if groups != 1:
             raise RuntimeError("Only support groups == 1")
         if dilation != 1:
@@ -911,7 +908,6 @@ def __init__(
             self.spatial_parallel_args = (1, 0, None, None, 0, False)
         else:
             self.spatial_parallel_args = spatial_parallel_args
-        return
 
     # Returns single callable that recomputes scale and bias for all frozen batch-norms.
     # This method must be called before cuda graphing.
diff --git a/apex/contrib/bottleneck/halo_exchangers.py b/apex/contrib/bottleneck/halo_exchangers.py
index eb0224c5b..8a9e82cac 100644
--- a/apex/contrib/bottleneck/halo_exchangers.py
+++ b/apex/contrib/bottleneck/halo_exchangers.py
@@ -1,13 +1,13 @@
-import torch
 import nccl_p2p_cuda as inc
 import peer_memory_cuda as pm
+import torch
 
 
 # Communication free halo exchanger.
 # NB! This halo exchanger does not exchange halos with neighbors as it should, it merely swaps the inputs
 # NB! This is only useful for performance testing.
 # NB! Do not use for actual production runs
-class HaloExchanger(object):
+class HaloExchanger:
     def __init__(self, ranks, rank_in_group):
         self.stream1 = torch.cuda.Stream()
         self.stream2 = torch.cuda.Stream()
@@ -27,7 +27,7 @@ def __init__(self, ranks, rank_in_group):
 
 class HaloExchangerNoComm(HaloExchanger):
     def __init__(self, ranks, rank_in_group):
-        super(HaloExchangerNoComm, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
 
     def left_right_halo_exchange(
         self,
@@ -45,7 +45,7 @@ def left_right_halo_exchange(
 
 class HaloExchangerAllGather(HaloExchanger):
     def __init__(self, ranks, rank_in_group, comm):
-        super(HaloExchangerAllGather, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
         # self.comm must be NCCL process_group created with torch.distributed.new_group(ranks=ranks)
         self.comm = comm
 
@@ -94,7 +94,7 @@ def left_right_halo_exchange(
 
 class HaloExchangerSendRecv(HaloExchanger):
     def __init__(self, ranks, rank_in_group):
-        super(HaloExchangerSendRecv, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
         nccl_id = inc.get_unique_nccl_id(1).cuda()
         torch.distributed.broadcast(nccl_id, 0)
         nccl_id = nccl_id.cpu()
@@ -145,7 +145,7 @@ def left_right_halo_exchange(
 
 class HaloExchangerPeer(HaloExchanger):
     def __init__(self, ranks, rank_in_group, peer_pool, explicit_nhwc, numSM=0):
-        super(HaloExchangerPeer, self).__init__(ranks, rank_in_group)
+        super().__init__(ranks, rank_in_group)
         self.diagnostics = False
         self.explicit_nhwc = explicit_nhwc
         self.numSM = numSM
diff --git a/apex/contrib/bottleneck/test.py b/apex/contrib/bottleneck/test.py
index d0e8703be..f57bdb4dd 100644
--- a/apex/contrib/bottleneck/test.py
+++ b/apex/contrib/bottleneck/test.py
@@ -29,7 +29,7 @@
         print("[DEBUG] ref dx :", d_grad.sum().item())
         # print wgrad. we don't need to reset since later cpp print before accumulation
         for i, w in enumerate(model.w_conv):
-            print("[DEBUG] ref wgrad{} :".format(i + 1), w.grad.sum().item())
+            print(f"[DEBUG] ref wgrad{i + 1} :", w.grad.sum().item())
 
     wgrads = []
     for w in model.w_conv:
@@ -56,7 +56,7 @@
     )
     for i, (w, wgrad) in enumerate(zip(model.w_conv, wgrads)):
         print(
-            "max error wgrad{}:".format(i + 1),
+            f"max error wgrad{i + 1}:",
             (wgrad - w.grad.float()).abs().max().item(),
             "max elem:",
             wgrad.abs().max().item(),
@@ -104,7 +104,7 @@
     )
     for i, (w, wgrad) in enumerate(zip(nhwc_model.w_conv, wgrads)):
         print(
-            "max error wgrad{}:".format(i + 1),
+            f"max error wgrad{i + 1}:",
             (wgrad - w.grad.float()).abs().max().item(),
             "max elem:",
             wgrad.abs().max().item(),
diff --git a/apex/contrib/clip_grad/clip_grad.py b/apex/contrib/clip_grad/clip_grad.py
index b34dc43bd..8ac27d98c 100644
--- a/apex/contrib/clip_grad/clip_grad.py
+++ b/apex/contrib/clip_grad/clip_grad.py
@@ -1,10 +1,12 @@
-from typing import Union, Iterable
+from collections.abc import Iterable
+from typing import Union
 
 import torch
 
 _kernel_import_succeeded = False
 try:
     import amp_C
+
     from apex.multi_tensor_apply import multi_tensor_applier
 
     _kernel_import_succeeded = True
diff --git a/apex/contrib/conv_bias_relu/__init__.py b/apex/contrib/conv_bias_relu/__init__.py
index ca18aa520..f4310fde9 100644
--- a/apex/contrib/conv_bias_relu/__init__.py
+++ b/apex/contrib/conv_bias_relu/__init__.py
@@ -1,6 +1,6 @@
 from .conv_bias_relu import (
-    ConvBiasReLU,
     ConvBias,
     ConvBiasMaskReLU,
+    ConvBiasReLU,
     ConvFrozenScaleBiasReLU,
 )
diff --git a/apex/contrib/conv_bias_relu/conv_bias_relu.py b/apex/contrib/conv_bias_relu/conv_bias_relu.py
index 533d6421f..f5d0c8389 100644
--- a/apex/contrib/conv_bias_relu/conv_bias_relu.py
+++ b/apex/contrib/conv_bias_relu/conv_bias_relu.py
@@ -1,7 +1,7 @@
+import fused_conv_bias_relu
 import torch
 
 from apex import check_cudnn_version_and_warn
-import fused_conv_bias_relu
 
 check_cudnn_version_and_warn(__name__, 8400)
 
diff --git a/apex/contrib/csrc/group_norm_v2/generate_gn_cuda_inst.py b/apex/contrib/csrc/group_norm_v2/generate_gn_cuda_inst.py
index b00141a77..7c3775424 100644
--- a/apex/contrib/csrc/group_norm_v2/generate_gn_cuda_inst.py
+++ b/apex/contrib/csrc/group_norm_v2/generate_gn_cuda_inst.py
@@ -1,6 +1,5 @@
 import pathlib
 
-
 hw_c_list = [
     (8 * 8, 1280),
     (8 * 8, 2560),
diff --git a/apex/contrib/cudnn_gbn/batch_norm.py b/apex/contrib/cudnn_gbn/batch_norm.py
index 8346b74aa..43eeabea4 100644
--- a/apex/contrib/cudnn_gbn/batch_norm.py
+++ b/apex/contrib/cudnn_gbn/batch_norm.py
@@ -1,10 +1,10 @@
+import cudnn_gbn_lib
+import peer_memory_cuda as pm
 import torch
-from torch.nn.modules.batchnorm import _BatchNorm
-from torch.nn import functional as F
 from torch import Tensor
-import peer_memory_cuda as pm
-import cudnn_gbn_lib
-from torch.cuda.amp import custom_fwd, custom_bwd
+from torch.cuda.amp import custom_bwd, custom_fwd
+from torch.nn import functional as F
+from torch.nn.modules.batchnorm import _BatchNorm
 
 
 class _GroupBatchNorm2d(torch.autograd.Function):
@@ -128,7 +128,7 @@ def __init__(
         affine=True,
         track_running_stats=True,
     ):
-        super(GroupBatchNorm2d, self).__init__(
+        super().__init__(
             num_features,
             eps=eps,
             momentum=momentum,
@@ -165,7 +165,7 @@ def get_peer_buffers(self, num_features):
 
     def _check_input_dim(self, input):
         if input.dim() != 4:
-            raise ValueError("expected 4D input (got {}D input)".format(input.dim()))
+            raise ValueError(f"expected 4D input (got {input.dim()}D input)")
 
     def _check_input_channels(self, input):
         if input.size(1) % 8 != 0:
diff --git a/apex/contrib/examples/gpu_direct_storage/benchmark_load.py b/apex/contrib/examples/gpu_direct_storage/benchmark_load.py
index d404aaac9..faf4eaa31 100644
--- a/apex/contrib/examples/gpu_direct_storage/benchmark_load.py
+++ b/apex/contrib/examples/gpu_direct_storage/benchmark_load.py
@@ -1,6 +1,8 @@
 import timeit
-import torch
+
 import apex.contrib.gpu_direct_storage as gds
+import torch
+
 
 def run_benchmark_torch_load():
     sizes = [2 ** i for i in range(16, 28)]
diff --git a/apex/contrib/examples/gpu_direct_storage/benchmark_save.py b/apex/contrib/examples/gpu_direct_storage/benchmark_save.py
index 36234b5a7..b6f64de09 100644
--- a/apex/contrib/examples/gpu_direct_storage/benchmark_save.py
+++ b/apex/contrib/examples/gpu_direct_storage/benchmark_save.py
@@ -1,7 +1,9 @@
 import os
 import timeit
-import torch
+
 import apex.contrib.gpu_direct_storage as gds
+import torch
+
 
 def run_benchmark(func):
     sizes = [2 ** i for i in range(16, 28)]
diff --git a/apex/contrib/examples/gpu_direct_storage/example_load.py b/apex/contrib/examples/gpu_direct_storage/example_load.py
index e471dbca1..7063ccc5f 100644
--- a/apex/contrib/examples/gpu_direct_storage/example_load.py
+++ b/apex/contrib/examples/gpu_direct_storage/example_load.py
@@ -1,5 +1,5 @@
-import torch
 import apex.contrib.gpu_direct_storage as gds
+import torch
 
 for size in [128, 1024, 8192]:
     x = torch.empty(size, device = "cuda")
diff --git a/apex/contrib/examples/gpu_direct_storage/example_save.py b/apex/contrib/examples/gpu_direct_storage/example_save.py
index 6b67d20c0..1af1fb475 100644
--- a/apex/contrib/examples/gpu_direct_storage/example_save.py
+++ b/apex/contrib/examples/gpu_direct_storage/example_save.py
@@ -1,5 +1,5 @@
-import torch
 import apex.contrib.gpu_direct_storage as gds
+import torch
 
 for size in [128, 1024, 8192]:
     x = torch.linspace(0, 1, size, device = "cuda")
diff --git a/apex/contrib/examples/nccl_allocator/allreduce.py b/apex/contrib/examples/nccl_allocator/allreduce.py
index 5089995af..eefcaf20f 100644
--- a/apex/contrib/examples/nccl_allocator/allreduce.py
+++ b/apex/contrib/examples/nccl_allocator/allreduce.py
@@ -1,7 +1,9 @@
 import os
+
 import torch
 import torch.distributed as dist
-import apex.contrib.nccl_allocator as nccl_allocator
+
+from apex.contrib import nccl_allocator
 
 assert os.getenv("WORLD_SIZE") is not None, "Please use: torchrun --nproc-per-node=8 allreduce.py"
 
diff --git a/apex/contrib/examples/nccl_allocator/cache.py b/apex/contrib/examples/nccl_allocator/cache.py
index 124f67062..699ebae62 100644
--- a/apex/contrib/examples/nccl_allocator/cache.py
+++ b/apex/contrib/examples/nccl_allocator/cache.py
@@ -1,7 +1,9 @@
 import torch
-import apex.contrib.nccl_allocator as nccl_allocator
 from pynvml.smi import nvidia_smi
 
+from apex.contrib import nccl_allocator
+
+
 def set_device(dev):
     import ctypes
     handle = ctypes.CDLL("libcudart.so")
diff --git a/apex/contrib/examples/nccl_allocator/change_cuda_allocator.py b/apex/contrib/examples/nccl_allocator/change_cuda_allocator.py
index 9311e1cfe..f4131a8c7 100644
--- a/apex/contrib/examples/nccl_allocator/change_cuda_allocator.py
+++ b/apex/contrib/examples/nccl_allocator/change_cuda_allocator.py
@@ -1,5 +1,6 @@
 import torch
-import apex.contrib.nccl_allocator as nccl_allocator
+
+from apex.contrib import nccl_allocator
 
 nccl_allocator.init()
 nrep = 6
diff --git a/apex/contrib/examples/nccl_allocator/toy_ddp.py b/apex/contrib/examples/nccl_allocator/toy_ddp.py
index 3da923bfa..47f868587 100644
--- a/apex/contrib/examples/nccl_allocator/toy_ddp.py
+++ b/apex/contrib/examples/nccl_allocator/toy_ddp.py
@@ -1,17 +1,17 @@
 import os
+
 import torch
-import torch.nn as nn
-import torch.optim as optim
 import torch.distributed as dist
+from torch import nn, optim
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-import apex.contrib.nccl_allocator as nccl_allocator
+from apex.contrib import nccl_allocator
 
 assert os.getenv("WORLD_SIZE") is not None, "Please use: torchrun --nproc-per-node=8 toy_ddp.py"
 
 class ToyModel(nn.Module):
     def __init__(self):
-        super(ToyModel, self).__init__()
+        super().__init__()
         self.net1 = nn.Linear(10, 10)
         self.relu = nn.ReLU()
         self.net2 = nn.Linear(10, 5)
diff --git a/apex/contrib/focal_loss/__init__.py b/apex/contrib/focal_loss/__init__.py
index 2a187d029..ca5081b12 100644
--- a/apex/contrib/focal_loss/__init__.py
+++ b/apex/contrib/focal_loss/__init__.py
@@ -1,6 +1,7 @@
 try:
-    import torch
     import focal_loss_cuda
+    import torch
+
     from .focal_loss import focal_loss
 
     del torch
diff --git a/apex/contrib/focal_loss/focal_loss.py b/apex/contrib/focal_loss/focal_loss.py
index 85c6f620e..2d76f3163 100644
--- a/apex/contrib/focal_loss/focal_loss.py
+++ b/apex/contrib/focal_loss/focal_loss.py
@@ -1,6 +1,5 @@
-import torch
-
 import focal_loss_cuda
+import torch
 
 
 class FocalLoss(torch.autograd.Function):
diff --git a/apex/contrib/group_norm/group_norm.py b/apex/contrib/group_norm/group_norm.py
index 998a220ff..56bb8a725 100644
--- a/apex/contrib/group_norm/group_norm.py
+++ b/apex/contrib/group_norm/group_norm.py
@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# coding: utf-8
 
 #
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
@@ -8,12 +7,12 @@
 
 import functools
 import os
-import torch
-import torch.nn.init as init
+
 import group_norm_cuda
 import group_norm_v2_cuda
-
+import torch
 from torch import Tensor
+from torch.nn import init
 from torch.nn.parameter import Parameter
 
 __all__ = ["GroupNorm"]
diff --git a/apex/contrib/groupbn/__init__.py b/apex/contrib/groupbn/__init__.py
index 4af4ac595..217fd64d4 100644
--- a/apex/contrib/groupbn/__init__.py
+++ b/apex/contrib/groupbn/__init__.py
@@ -1,6 +1,7 @@
 try:
-    import torch
     import bnp
+    import torch
+
     from .batch_norm import BatchNorm2d_NHWC
 
     del torch
diff --git a/apex/contrib/groupbn/batch_norm.py b/apex/contrib/groupbn/batch_norm.py
index d5d71cd0f..ddcd2651e 100644
--- a/apex/contrib/groupbn/batch_norm.py
+++ b/apex/contrib/groupbn/batch_norm.py
@@ -1,9 +1,8 @@
-import torch
+import bnp
 import numpy as np
+import torch
 from torch.nn.modules.batchnorm import _BatchNorm
 
-import bnp
-
 
 class bn_NHWC_impl(torch.autograd.Function):
     @staticmethod
@@ -298,7 +297,7 @@ def __init__(
         cta_launch_margin=12,
         multi_stream=False,
     ):
-        super(BatchNorm2d_NHWC, self).__init__(num_features)
+        super().__init__(num_features)
 
         self.fuse_relu = fuse_relu
         self.multi_stream = multi_stream
diff --git a/apex/contrib/index_mul_2d/index_mul_2d.py b/apex/contrib/index_mul_2d/index_mul_2d.py
index ab628b05d..89f2d9692 100644
--- a/apex/contrib/index_mul_2d/index_mul_2d.py
+++ b/apex/contrib/index_mul_2d/index_mul_2d.py
@@ -1,6 +1,5 @@
-import torch
-
 import fused_index_mul_2d
+import torch
 
 
 class IndexMul2d_(torch.autograd.Function):
diff --git a/apex/contrib/layer_norm/layer_norm.py b/apex/contrib/layer_norm/layer_norm.py
index ef134667d..4da056657 100644
--- a/apex/contrib/layer_norm/layer_norm.py
+++ b/apex/contrib/layer_norm/layer_norm.py
@@ -1,8 +1,8 @@
+import fast_layer_norm
 import torch
 from torch.nn import init
 
 from apex._autocast_utils import _cast_if_autocast_enabled
-import fast_layer_norm
 
 
 class FastLayerNormFN(torch.autograd.Function):
diff --git a/apex/contrib/nccl_allocator/nccl_allocator.py b/apex/contrib/nccl_allocator/nccl_allocator.py
index 8232021e0..f68e387c9 100644
--- a/apex/contrib/nccl_allocator/nccl_allocator.py
+++ b/apex/contrib/nccl_allocator/nccl_allocator.py
@@ -1,11 +1,10 @@
 import os
-import torch
-import _apex_nccl_allocator
-
 from contextlib import nullcontext
 
+import _apex_nccl_allocator
+import torch
 
-__all__ = ["init", "nccl_mem", "create_nccl_mem_pool"]
+__all__ = ["create_nccl_mem_pool", "init", "nccl_mem"]
 
 
 def get_func_args(func):
diff --git a/apex/contrib/openfold_triton/__init__.py b/apex/contrib/openfold_triton/__init__.py
index 771c71e14..84b77b70e 100644
--- a/apex/contrib/openfold_triton/__init__.py
+++ b/apex/contrib/openfold_triton/__init__.py
@@ -30,16 +30,16 @@
 )
 
 __all__ = (
-    "LayerNormSmallShapeOptImpl",
-    "sync_triton_auto_tune_cache_across_gpus",
-    "CanSchTriMHA",
-    "AttnTri",
     "AttnBiasJIT",
     "AttnNoBiasJIT",
+    "AttnTri",
+    "CanSchTriMHA",
+    "LayerNormSmallShapeOptImpl",
+    "sync_triton_auto_tune_cache_across_gpus",
 )
 
 
-def _get_tuneable_triton_func_name(f: Union[Autotuner, Heuristics, JITFunction]) -> str:
+def _get_tuneable_triton_func_name(f: Autotuner | Heuristics | JITFunction) -> str:
     if isinstance(f, JITFunction):
         return f.__name__
     else:
diff --git a/apex/contrib/openfold_triton/_layer_norm_forward_kernels.py b/apex/contrib/openfold_triton/_layer_norm_forward_kernels.py
index 31d66a6a1..6512bf2a9 100644
--- a/apex/contrib/openfold_triton/_layer_norm_forward_kernels.py
+++ b/apex/contrib/openfold_triton/_layer_norm_forward_kernels.py
@@ -1,9 +1,8 @@
 # © 2023 NVIDIA CORPORATION & AFFILIATES
 
-from packaging.version import Version
-
 import triton
 import triton.language as tl
+from packaging.version import Version
 from triton import Config
 
 if Version("2.0.0") < Version(triton.__version__):
diff --git a/apex/contrib/openfold_triton/_mha_kernel.py b/apex/contrib/openfold_triton/_mha_kernel.py
index 11f44c884..b27c56220 100644
--- a/apex/contrib/openfold_triton/_mha_kernel.py
+++ b/apex/contrib/openfold_triton/_mha_kernel.py
@@ -461,7 +461,7 @@ def _bwd_kernel(
         Mask += off_b * stride_mz + off_h * stride_mh
 
     num_block_n = tl.cdiv(N_CTX, BLOCK_N)
-    for start_n in range(0, num_block_n):
+    for start_n in range(num_block_n):
         # lo = start_n * BLOCK_M
         lo = 0
         # initialize row/col offsets
diff --git a/apex/contrib/openfold_triton/fused_adam_swa.py b/apex/contrib/openfold_triton/fused_adam_swa.py
index 11a654812..e41ce9500 100644
--- a/apex/contrib/openfold_triton/fused_adam_swa.py
+++ b/apex/contrib/openfold_triton/fused_adam_swa.py
@@ -3,14 +3,15 @@
 from __future__ import annotations
 
 from collections import defaultdict
+from collections.abc import Callable
 from enum import Enum, unique
 from itertools import chain
-from typing import Callable, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
-import torch.nn as nn
 import triton
 import triton.language as tl
+from torch import nn
 from torch.optim import Adam, Optimizer
 
 # The most common parameter size in open-fold.
@@ -209,13 +210,13 @@ def _multi_tensor_adam_swa(
 class FusedAdamSWA(Optimizer):
     def __init__(
         self,
-        params: List[nn.Parameter],
-        compute_params: List[nn.Parameter],
-        swa_params: List[nn.Parameter],
+        params: list[nn.Parameter],
+        compute_params: list[nn.Parameter],
+        swa_params: list[nn.Parameter],
         swa_decay_rate: float,
         lr: float = 1e-3,
         bias_correction: bool = True,
-        betas: Tuple[float, float] = (0.9, 0.999),
+        betas: tuple[float, float] = (0.9, 0.999),
         eps: float = 1e-8,
         adam_math_mode: AdamMathType = AdamMathType.PyTorchAdam,
         weight_decay: float = 0.0,
@@ -371,8 +372,8 @@ def _build_pointer_buffers(self):
 
     def step(
         self,
-        closure: Optional[Callable[[], torch.Tensor]] = None,
-        grad_clip_scale: Optional[Union[torch.Tensor, float]] = None,
+        closure: Callable[[], torch.Tensor] | None = None,
+        grad_clip_scale: torch.Tensor | float | None = None,
     ):
         if not self._pointer_buffers_initialized:
             self._build_pointer_buffers()
@@ -460,9 +461,9 @@ def step(
     def from_optim(
         cls,
         adam_optimizer: Adam,
-        fp32_params: List[nn.Parameter],
-        bf16_params: List[nn.Parameter],
-        swa_params: List[nn.Parameter],
+        fp32_params: list[nn.Parameter],
+        bf16_params: list[nn.Parameter],
+        swa_params: list[nn.Parameter],
         swa_decay_rate: float,
     ) -> FusedAdamSWA:
         assert len(adam_optimizer.param_groups) == 1
diff --git a/apex/contrib/openfold_triton/mha.py b/apex/contrib/openfold_triton/mha.py
index 9065b6ca8..07b4a5b07 100644
--- a/apex/contrib/openfold_triton/mha.py
+++ b/apex/contrib/openfold_triton/mha.py
@@ -17,7 +17,7 @@
 _TRI_MHA_ENABLED = False
 
 
-def is_enabled() -> Optional[bool]:
+def is_enabled() -> bool | None:
     global _TRI_MHA_ENABLED
     return _TRI_MHA_ENABLED
 
@@ -90,38 +90,34 @@ def CanSchTriMHA(in_shape, has_bias=True, inf=1e9, training=True):
 def schedule_triton_mha(in_shape, fwd=True):
     # default
     ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 0]
-    if in_shape == [256, 4, 256, 16]:
-        ret = [64, 32, 2, 4] if fwd else [64, 64, 4, 0]
-    elif in_shape == [128, 4, 256, 16]:
-        ret = [64, 32, 2, 4] if fwd else [64, 64, 4, 0]
-    elif in_shape == [64, 4, 256, 16]:
-        ret = [64, 32, 2, 4] if fwd else [64, 64, 4, 0]
-    elif in_shape == [32, 4, 256, 16]:
+    if (
+        in_shape == [256, 4, 256, 16]
+        or in_shape == [128, 4, 256, 16]
+        or in_shape == [64, 4, 256, 16]
+        or in_shape == [32, 4, 256, 16]
+    ):
         ret = [64, 32, 2, 4] if fwd else [64, 64, 4, 0]
     # [*, 8, 256, 32]
-    elif in_shape == [128, 8, 256, 32]:  # DAP1
-        ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
-    elif in_shape == [64, 8, 256, 32]:  # DAP2
-        ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
-    elif in_shape == [32, 8, 256, 32]:  # DAP4
-        ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
-    elif in_shape == [16, 8, 256, 32]:  # DAP8
+    elif (
+        in_shape == [128, 8, 256, 32]
+        or in_shape == [64, 8, 256, 32]
+        or in_shape == [32, 8, 256, 32]
+        or in_shape == [16, 8, 256, 32]
+    ):  # DAP1
         ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
     # [*, 8, 128, 32]
     elif in_shape == [256, 8, 128, 32]:  # DAP1
         ret = [64, 64, 4, 3] if fwd else [128, 64, 4, 1]
-    elif in_shape == [128, 8, 128, 32]:  # DAP2
-        ret = [128, 64, 4, 2] if fwd else [64, 64, 2, 0]
-    elif in_shape == [64, 8, 128, 32]:  # DAP4
-        ret = [128, 64, 4, 2] if fwd else [64, 64, 2, 0]
-    elif in_shape == [32, 8, 128, 32]:  # DAP8
+    elif (
+        in_shape == [128, 8, 128, 32]
+        or in_shape == [64, 8, 128, 32]
+        or in_shape == [32, 8, 128, 32]
+    ):  # DAP2
         ret = [128, 64, 4, 2] if fwd else [64, 64, 2, 0]
     # [*, 4, 256, 32]
     elif in_shape == [256, 4, 256, 32]:  # DAP1
         ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 0]
-    elif in_shape == [128, 4, 256, 32]:  # DAP2
-        ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
-    elif in_shape == [64, 4, 256, 32]:  # DAP4
+    elif in_shape == [128, 4, 256, 32] or in_shape == [64, 4, 256, 32]:  # DAP2
         ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 1]
     elif in_shape == [32, 4, 256, 32]:  # DAP8
         ret = [64, 32, 2, 3] if fwd else [128, 64, 8, 0]
@@ -398,7 +394,7 @@ def _attention_bias(
     key: torch.Tensor,
     value: torch.Tensor,
     mask: torch.Tensor,
-    bias: Optional[torch.Tensor],
+    bias: torch.Tensor | None,
     inf: float,
 ) -> torch.Tensor:
     # query:  [*, num_heads, Q, c_hidden]
diff --git a/apex/contrib/optimizers/distributed_fused_adam.py b/apex/contrib/optimizers/distributed_fused_adam.py
index b1f6b58fc..e3dae9a07 100644
--- a/apex/contrib/optimizers/distributed_fused_adam.py
+++ b/apex/contrib/optimizers/distributed_fused_adam.py
@@ -1,36 +1,36 @@
 import collections
 import contextlib
-from dataclasses import dataclass
 import enum
 import inspect
 import io
 import itertools
 import threading
+import warnings
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
 from typing import (
     Any,
-    Callable,
     Dict,
-    Iterable,
     List,
     Optional,
     Set,
     Tuple,
     Union,
 )
-import warnings
 
 import torch
 from torch.distributed.distributed_c10d import _get_default_group
 
 try:
-    import apex.contrib.nccl_allocator as nccl_allocator
+    from apex.contrib import nccl_allocator
 except ImportError:
     nccl_allocator = None
 
-from apex.multi_tensor_apply import multi_tensor_applier
 import amp_C
 import distributed_adam_cuda
 
+from apex.multi_tensor_apply import multi_tensor_applier
+
 # Fallback to private functions if using PyTorch <1.13.0
 try:
     from torch.distributed.distributed_c10d import get_global_rank
@@ -73,7 +73,7 @@ def _coalescing_manager(group, device, reqs):
 
     class _CoalescingManager:
         def __init__(self):
-            self.works: List[torch.distributed.Work] = []
+            self.works: list[torch.distributed.Work] = []
 
         def append(self, work: torch.distributed.Work) -> None:
             if work:
@@ -85,8 +85,8 @@ def wait(self) -> None:
 
     @contextlib.contextmanager
     def _coalescing_manager(
-        group: Optional[torch.distributed.ProcessGroup] = None,
-        device: Optional[torch.device] = None,
+        group: torch.distributed.ProcessGroup | None = None,
+        device: torch.device | None = None,
         async_ops: bool = False,
     ) -> contextlib.AbstractContextManager:
         assert device is not None
@@ -120,7 +120,6 @@ def _coalescing_manager_append_work(
         communication.
 
         """
-        pass
 
 
 # Import optional CUDA kernels
@@ -166,9 +165,9 @@ def _devices_match(device1: torch.device, device2: torch.device) -> bool:
 
 
 def _multi_tensor_copy(
-    buffers_in: List[torch.Tensor],
-    buffers_out: List[torch.Tensor],
-    dummy_overflow_buf: Optional[torch.Tensor] = None,
+    buffers_in: list[torch.Tensor],
+    buffers_out: list[torch.Tensor],
+    dummy_overflow_buf: torch.Tensor | None = None,
 ) -> None:
     """Copy between corresponding buffers
 
@@ -268,7 +267,7 @@ def _bf16_rem_to_fp32(
 
 
 class DistributedFusedAdam(torch.optim.Optimizer):
-    """Adam optimizer with ZeRO algorithm.
+    r"""Adam optimizer with ZeRO algorithm.
 
     Currently GPU-only. Requires Apex to be installed via
     ``python setup.py install --cuda_ext --cpp_ext --distributed_adam --deprecated_fused_adam``.
@@ -401,17 +400,17 @@ class ParameterFragment:
         # Bucket index
         bucket_id: int
         # Range within flattened parameter buffer
-        param_range: Tuple[int, int]
+        param_range: tuple[int, int]
         # Range within bucket
-        bucket_range: Tuple[int, int]
+        bucket_range: tuple[int, int]
         # Whether fragment is in local shard of bucket
         in_local_shard: bool
         # Range within local shard
-        shard_range: Optional[Tuple[int, int]]
+        shard_range: tuple[int, int] | None
         # Range of local fragment shard within bucket
-        shard_bucket_range: Optional[Tuple[int, int]]
+        shard_bucket_range: tuple[int, int] | None
         # Range of local fragment shard within parameter
-        shard_param_range: Optional[Tuple[int, int]]
+        shard_param_range: tuple[int, int] | None
 
     class StateBucket:
         """Optimizer state for a bucket"""
@@ -445,9 +444,9 @@ def __init__(
             # Offset to bucket in contiguous buffers
             self.contiguous_buffer_offset: int = contiguous_buffer_offset
             # Buffer ranges corresponding to parameter fragments
-            self.fragments: List[ParameterFragment] = []
+            self.fragments: list[ParameterFragment] = []
             # Local shard of parameters
-            self.params_shard: Optional[torch.Tensor] = None
+            self.params_shard: torch.Tensor | None = None
             if store_params:
                 self.params_shard = torch.zeros(
                     [shard_size],
@@ -455,7 +454,7 @@ def __init__(
                     device=device,
                 )
             # Local shard of parameter remainders
-            self.param_remainders_shard: Optional[torch.Tensor] = None
+            self.param_remainders_shard: torch.Tensor | None = None
             if store_param_remainders:
                 self.param_remainders_shard = torch.zeros(
                     [shard_size],
@@ -475,7 +474,7 @@ def __init__(
                 device=device,
             )
 
-        def dtypes(self) -> Tuple[torch.dtype, torch.dtype, torch.dtype]:
+        def dtypes(self) -> tuple[torch.dtype, torch.dtype, torch.dtype]:
             """Datatypes for the bucket's compute and communication"""
             return (
                 self.dtype,
@@ -500,15 +499,15 @@ class GradientBucket:
 
         def __init__(self):
             # Local shard of gradients
-            self.grads_shard: Optional[torch.Tensor] = None
+            self.grads_shard: torch.Tensor | None = None
             # Local contribution to gradients
-            self.grads_bucket: Optional[torch.Tensor] = None
+            self.grads_bucket: torch.Tensor | None = None
             # Buffer for gradient reduce-scatter
-            self.sync_grads_shard: Optional[torch.Tensor] = None
+            self.sync_grads_shard: torch.Tensor | None = None
             # Status of gradients
             self.status: GradientStatus = DistributedFusedAdam.GradientStatus.READY
             # Params that have generated grads
-            self.grads_generated: Set[torch.nn.Parameter] = set()
+            self.grads_generated: set[torch.nn.Parameter] = set()
 
     class ParameterStatus(enum.Enum):
         """Status of parameters within a bucket"""
@@ -525,13 +524,13 @@ class ParameterBucket:
 
         def __init__(self):
             # Local shard of parameters
-            self.params_shard: Optional[torch.Tensor] = None
+            self.params_shard: torch.Tensor | None = None
             # Gathered parameter values
-            self.params_bucket: Optional[torch.Tensor] = None
+            self.params_bucket: torch.Tensor | None = None
             # Status of parameters
             self.status: ParameterStatus = DistributedFusedAdam.ParameterStatus.SHARDED
             # Params that have been updated
-            self.params_updated: Set[torch.nn.Parameter] = set()
+            self.params_updated: set[torch.nn.Parameter] = set()
 
     # Enable custom logic for AMP grad scaling
     _step_supports_amp_scaling: bool = True
@@ -539,21 +538,21 @@ def __init__(self):
 
     def __init__(
         self,
-        params: Union[Iterable[torch.nn.Parameter], Iterable[dict]],
+        params: Iterable[torch.nn.Parameter] | Iterable[dict],
         lr: float = 1e-3,
         bias_correction: bool = True,
-        betas: Tuple[float, float] = (0.9, 0.999),
+        betas: tuple[float, float] = (0.9, 0.999),
         eps: float = 1e-8,
         adam_w_mode: bool = True,
         weight_decay: float = 0.0,
         amsgrad: bool = False,
         dtype: torch.dtype = torch.float32,
-        grad_sync_dtype: Optional[torch.dtype] = None,
-        param_sync_dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = "cuda",
-        process_group: Optional[torch.distributed.ProcessGroup] = None,
-        distributed_process_group: Optional[torch.distributed.ProcessGroup] = None,
-        redundant_process_group: Optional[torch.distributed.ProcessGroup] = None,
+        grad_sync_dtype: torch.dtype | None = None,
+        param_sync_dtype: torch.dtype | None = None,
+        device: torch.device | None = "cuda",
+        process_group: torch.distributed.ProcessGroup | None = None,
+        distributed_process_group: torch.distributed.ProcessGroup | None = None,
+        redundant_process_group: torch.distributed.ProcessGroup | None = None,
         average_grad_sync: bool = True,
         overlap_grad_sync: bool = True,
         overlap_param_sync: bool = False,
@@ -631,7 +630,7 @@ def __init__(
         self.distributed_process_group: torch.distributed.ProcessGroup = (
             self.process_group if distributed_process_group is None else distributed_process_group
         )
-        self.redundant_process_group: Optional[torch.distributed.ProcessGroup] = (
+        self.redundant_process_group: torch.distributed.ProcessGroup | None = (
             redundant_process_group
         )
         self.process_group_size: int = torch.distributed.get_world_size(self.process_group)
@@ -726,17 +725,17 @@ def __init__(
         self.default_shard_size: int = shard_size
 
         # Optimizer state
-        self.state["buckets"]: List[StateBucket] = []
+        self.state["buckets"]: list[StateBucket] = []
         self.state["step"]: torch.Tensor | int = (
             torch.tensor([0], dtype=torch.int, device=self.device) if self.capturable else 0
         )
 
         # Gradient state
-        self._grads_buckets: Dict[int, GradientBucket] = collections.defaultdict(
+        self._grads_buckets: dict[int, GradientBucket] = collections.defaultdict(
             self.GradientBucket
         )
         # Param state
-        self._params_buckets: Dict[int, ParameterBucket] = collections.OrderedDict()
+        self._params_buckets: dict[int, ParameterBucket] = collections.OrderedDict()
 
         # Whether to allocate contiguous buffers for parameters
         self.contiguous_param_buffer: bool = contiguous_param_buffer
@@ -745,9 +744,9 @@ def __init__(
         # Whether to use NCCL User Buffer
         self.nccl_ub: bool = nccl_ub
         # Contiguous buffers for parameters
-        self._param_buffers: Dict[Tuple[torch.dtype, torch.dtype, torch.dtype], torch.Tensor] = {}
+        self._param_buffers: dict[tuple[torch.dtype, torch.dtype, torch.dtype], torch.Tensor] = {}
         # Contiguous buffers for gradients
-        self._grad_buffers: Dict[Tuple[torch.dtype, torch.dtype, torch.dtype], torch.Tensor] = {}
+        self._grad_buffers: dict[tuple[torch.dtype, torch.dtype, torch.dtype], torch.Tensor] = {}
         # Output buffer for gradient shards, only required for NCCL user buffer
         if self.nccl_ub:
             if not nccl_allocator:
@@ -755,16 +754,16 @@ def __init__(
             elif not self.contiguous_grad_buffer:
                 raise RuntimeError("NCCL user buffers require contiguous grad buffers")
             else:
-                self._shard_grad_buffers: Dict[
-                    Tuple[torch.dtype, torch.dtype, torch.dtype], torch.Tensor
+                self._shard_grad_buffers: dict[
+                    tuple[torch.dtype, torch.dtype, torch.dtype], torch.Tensor
                 ] = {}
 
         # Side streams for state dict communication
-        self._pipeline_streams: List[torch.cuda.Stream] = [
+        self._pipeline_streams: list[torch.cuda.Stream] = [
             torch.cuda.Stream() for _ in range(self.pipeline_size)
         ]
         # Side streams for gradients and parameters communication
-        self._comm_streams: List[torch.cuda.Stream] = [
+        self._comm_streams: list[torch.cuda.Stream] = [
             torch.cuda.Stream() for _ in range(self.pipeline_size)
         ]
         self._last_comm_stream_id: int = -1
@@ -776,7 +775,7 @@ def __init__(
         )
         # Norm of parameter gradients. Used for gradient clipping and
         # gradient scaler.
-        self._grad_norm: Optional[torch.Tensor] = None
+        self._grad_norm: torch.Tensor | None = None
 
         # Dummy flag for multi-tensor kernels
         # Note: Apex multi-tensor kernels have a noop_flag argument
@@ -1198,7 +1197,7 @@ def parameters(self) -> Iterable[torch.nn.Parameter]:
 
     def parameter(
         self,
-        *args: Union[int, ParameterFragment],
+        *args: int | ParameterFragment,
     ) -> torch.nn.Parameter:
         """Get optimizer parameter
 
@@ -1227,10 +1226,10 @@ def parameter(
 
     def init_params(
         self,
-        params: Optional[Iterable[torch.nn.Parameter]] = None,
-        dtype: Optional[torch.dtype] = None,
-        grad_sync_dtype: Optional[torch.dtype] = None,
-        param_sync_dtype: Optional[torch.dtype] = None,
+        params: Iterable[torch.nn.Parameter] | None = None,
+        dtype: torch.dtype | None = None,
+        grad_sync_dtype: torch.dtype | None = None,
+        param_sync_dtype: torch.dtype | None = None,
     ) -> None:
         """Initialize optimizer state for parameters
 
@@ -1275,9 +1274,9 @@ def init_params(
     def init_params_bucket(
         self,
         params: Iterable[torch.nn.Parameter],
-        dtype: Optional[torch.dtype] = None,
-        grad_sync_dtype: Optional[torch.dtype] = None,
-        param_sync_dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
+        grad_sync_dtype: torch.dtype | None = None,
+        param_sync_dtype: torch.dtype | None = None,
     ) -> None:
         """Initialize optimizer state for parameters in one effective bucket
 
@@ -1349,9 +1348,9 @@ def _init_param_state(
         param: torch.nn.Parameter,
         param_group_id: int,
         param_id: int,
-        dtype: Optional[torch.dtype] = None,
-        grad_sync_dtype: Optional[torch.dtype] = None,
-        param_sync_dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
+        grad_sync_dtype: torch.dtype | None = None,
+        param_sync_dtype: torch.dtype | None = None,
     ) -> None:
         """Initialize optimizer state for a parameter"""
 
@@ -1667,7 +1666,7 @@ def _grad_copy(self, param: torch.nn.Parameter) -> None:
 
     def _param_copy(
         self,
-        params: Union[torch.nn.Parameter, Iterable[torch.nn.Parameter]],
+        params: torch.nn.Parameter | Iterable[torch.nn.Parameter],
     ) -> None:
         """Update parameters with values from parameter buckets
 
@@ -1826,7 +1825,7 @@ def _force_bucket_grad_sync(self) -> None:
 
     def _try_start_bucket_grad_sync(
         self,
-        params: Optional[Iterable[torch.nn.Parameter]] = None,
+        params: Iterable[torch.nn.Parameter] | None = None,
         ignore_last_bucket: bool = False,
     ) -> None:
         """Attempt to launch gradient synchronization
@@ -1874,7 +1873,7 @@ def _try_start_bucket_grad_sync(
         if filled_buckets:
             self._start_bucket_grad_sync(filled_buckets)
 
-    def _start_bucket_grad_sync(self, buckets: List[GradientBucket]) -> None:
+    def _start_bucket_grad_sync(self, buckets: list[GradientBucket]) -> None:
         """Synchronize gradient buckets
 
         Gradient synchronization is asynchronous. Involves
@@ -2029,7 +2028,7 @@ def _try_start_bucket_param_sync(
         if buckets:
             self._start_bucket_param_sync(buckets)
 
-    def _start_bucket_param_sync(self, buckets: List[ParameterBucket]) -> None:
+    def _start_bucket_param_sync(self, buckets: list[ParameterBucket]) -> None:
         """Synchronize parameter buckets
 
         Parameter synchronization is asynchronous. Involves all-gather
@@ -2141,7 +2140,7 @@ def param_sync(self) -> None:
             self._param_copy(self.parameters())
         else:
             while self._params_buckets:
-                bucket_id, bucket = next(iter((self._params_buckets.items())))
+                bucket_id, bucket = next(iter(self._params_buckets.items()))
                 for fragment in reversed(self.state["buckets"][bucket_id].fragments):
                     self._param_copy(self.parameter(fragment))
         self._params_buckets.clear()
@@ -2149,7 +2148,7 @@ def param_sync(self) -> None:
     @torch.no_grad()
     def _local_grad_norm(
         self,
-        parameters: Optional[Iterable[torch.nn.Parameter]] = None,
+        parameters: Iterable[torch.nn.Parameter] | None = None,
         norm_type: float = 2.0,
     ) -> torch.Tensor:
         """Local contribution to parameter gradient norm
@@ -2235,7 +2234,7 @@ def _local_grad_norm(
 
     def grad_norm(
         self,
-        parameters: Optional[Iterable[torch.nn.Parameter]] = None,
+        parameters: Iterable[torch.nn.Parameter] | None = None,
         norm_type: float = 2.0,
         force: bool = False,
     ) -> torch.Tensor:
@@ -2275,7 +2274,7 @@ def grad_norm(
     def clip_grad_norm(
         self,
         max_norm: float,
-        parameters: Optional[Iterable[torch.nn.Parameter]] = None,
+        parameters: Iterable[torch.nn.Parameter] | None = None,
         norm_type: float = 2.0,
     ) -> torch.Tensor:
         """Clips gradient norm of parameters in optimizer
@@ -2306,9 +2305,9 @@ def clip_grad_norm(
     @torch.no_grad
     def unscale_grads(
         self,
-        *args: Union[Optional[torch.Tensor], Any],
-        inv_scale: Optional[torch.Tensor] = None,
-        grad_scaler: Optional[torch.cuda.amp.GradScaler] = None,
+        *args: torch.Tensor | None | Any,
+        inv_scale: torch.Tensor | None = None,
+        grad_scaler: torch.cuda.amp.GradScaler | None = None,
     ) -> None:
         """Custom unscale function for use by AMP gradient scaler
 
@@ -2367,9 +2366,9 @@ def unscale_grads(
 
     def step(
         self,
-        closure: Optional[Callable] = None,
+        closure: Callable | None = None,
         *,
-        grad_scaler: Optional[torch.cuda.amp.GradScaler] = None,
+        grad_scaler: torch.cuda.amp.GradScaler | None = None,
     ):
         """Apply Adam optimizer step
 
@@ -2502,7 +2501,7 @@ def step(
 
         return loss
 
-    def _local_step(self, bucket_ids: List[int]) -> None:
+    def _local_step(self, bucket_ids: list[int]) -> None:
         """Apply optimizer step to local shard of parameter buckets
 
         Arguments:
@@ -2610,7 +2609,7 @@ def _local_step(self, bucket_ids: List[int]) -> None:
 
     def _local_step_with_param_remainders(
         self,
-        bucket_ids: List[int],
+        bucket_ids: list[int],
     ) -> None:
         """Apply optimizer step to local shard of parameter bucket
 
@@ -2693,7 +2692,7 @@ def _local_step_with_param_remainders(
     @torch.no_grad()
     def _local_step_with_scaled_states(
         self,
-        bucket_ids: List[int],
+        bucket_ids: list[int],
     ) -> None:
         for bucket_id in bucket_ids:
             state_bucket = self.state["buckets"][bucket_id]
@@ -2776,7 +2775,7 @@ def _local_step_with_scaled_states(
     @torch.no_grad()
     def _check_params_shard_dtypes(
         self,
-        params_buckets: Dict[int, ParameterBucket],
+        params_buckets: dict[int, ParameterBucket],
     ) -> None:
         """Make sure local shards of parameters are in expected datatypes
 
@@ -2862,9 +2861,9 @@ def _apply_state_scale(
     def state_dict(
         self,
         *,
-        state_dict_format: Optional[int] = None,
-        gather_on_root: Optional[bool] = None,
-    ) -> Optional[dict]:
+        state_dict_format: int | None = None,
+        gather_on_root: bool | None = None,
+    ) -> dict | None:
         """Get dictionary containing optimizer state
 
         All ranks in the process group must call this function since
@@ -2904,7 +2903,7 @@ def state_dict(
 
         return state_dict
 
-    def _state_dict_v1(self, gather_on_root: bool = True) -> Optional[dict]:
+    def _state_dict_v1(self, gather_on_root: bool = True) -> dict | None:
         """Get dictionary containing optimizer state (deprecated v1 format)
 
         Default behavior is to perform communication so that the
@@ -3056,7 +3055,7 @@ def _state_dict_v1(self, gather_on_root: bool = True) -> Optional[dict]:
             return None
 
     @torch.no_grad()
-    def _state_dict_v2(self) -> Optional[dict]:
+    def _state_dict_v2(self) -> dict | None:
         """Get dictionary containing optimizer state (default v2 format)
 
         All ranks in the process group must call this function since
diff --git a/apex/contrib/optimizers/distributed_fused_lamb.py b/apex/contrib/optimizers/distributed_fused_lamb.py
index 93b964fbb..29a95e872 100644
--- a/apex/contrib/optimizers/distributed_fused_lamb.py
+++ b/apex/contrib/optimizers/distributed_fused_lamb.py
@@ -1,12 +1,13 @@
-import os
-import inspect
-import torch
 import importlib
-import amp_C
-from apex.multi_tensor_apply import multi_tensor_applier
+import inspect
+import os
 
+import amp_C
+import torch
 import torch.distributed.distributed_c10d as c10d
 
+from apex.multi_tensor_apply import multi_tensor_applier
+
 # Fallback to private fields if using older PyTorch version
 try:
     import torch.distributed.distributed_c10d.get_process_group_ranks
@@ -83,7 +84,7 @@ class DistributedFusedLAMB(torch.optim.Optimizer):
         https://openreview.net/forum?id=ryQu7f-RZ
     """
 
-    class AtomicCounter(object):
+    class AtomicCounter:
         def __init__(self):
             self.value = 0
             self.order = []
@@ -137,7 +138,7 @@ def __init__(
             max_grad_norm=max_grad_norm,
         )
 
-        super(DistributedFusedLAMB, self).__init__(params, defaults)
+        super().__init__(params, defaults)
 
         global fused_adam_cuda, distributed_lamb_cuda
         fused_adam_cuda = importlib.import_module("fused_adam_cuda")
@@ -222,7 +223,7 @@ def __init__(
 
             self._ar_pg = []
             # consider all the ranks
-            ranks = list(range(0, self._world_size))
+            ranks = list(range(self._world_size))
             for i in range(self._num_ar_pg):
                 if self._verbose:
                     print(f"creating new AR group {i}: {ranks}")
@@ -665,12 +666,8 @@ def _get_flat_view(param):
                     ):
                         flat_grad_start = grads_info["param_offset"]
                         flat_grad_end = flat_grad_start + grads_info["param_grads_size"]
-                        clipped_start = (lambda a, b: a if a > b else b)(
-                            flat_grad_start, flat_shard_start
-                        )
-                        clipped_end = (lambda a, b: a if a < b else b)(
-                            flat_grad_end, flat_shard_end
-                        )
+                        clipped_start = (lambda a, b: max(b, a))(flat_grad_start, flat_shard_start)
+                        clipped_end = (lambda a, b: min(b, a))(flat_grad_end, flat_shard_end)
                         if clipped_start < clipped_end:
                             grad_offset = clipped_start - flat_grad_start
                             grad_length = clipped_end - clipped_start
diff --git a/apex/contrib/optimizers/fp16_optimizer.py b/apex/contrib/optimizers/fp16_optimizer.py
index 856a181dc..bf0af34e7 100755
--- a/apex/contrib/optimizers/fp16_optimizer.py
+++ b/apex/contrib/optimizers/fp16_optimizer.py
@@ -1,8 +1,9 @@
 import torch
+
 from apex.multi_tensor_apply import multi_tensor_applier
 
 
-class FP16_Optimizer(object):
+class FP16_Optimizer:
     """
     :class:`FP16_Optimizer` A cutdown version of apex.fp16_utils.FP16_Optimizer.
     Designed only to wrap apex.contrib.optimizers.FusedAdam, FusedSGD.
@@ -37,7 +38,7 @@ def __init__(
         # 1. maintain same user API from apex.fp16_utils
         # 2. keep common stuff here in case we need to add new fused optimizer later
 
-        if not torch.cuda.is_available:
+        if not torch.cuda.is_available():
             raise SystemError("Cannot use fp16 without CUDA.")
         self.optimizer = init_optimizer
 
@@ -112,7 +113,7 @@ def step(self, closure=None):
         self.overflow_buf.zero_()
         for fp16_grad in fp16_grads:
             if len(fp16_grad) > 0:
-                norm, norm_per_tensor = multi_tensor_applier(
+                norm, _norm_per_tensor = multi_tensor_applier(
                     self.multi_tensor_l2norm, self.overflow_buf, [fp16_grad], True
                 )
                 norm_groups.append(norm)
@@ -160,7 +161,6 @@ def _update_scale(self, skip):
                 print("\nGrad overflow on iteration", self.cur_iter)
                 print("Using static loss scale of", self.cur_scale)
         self.cur_iter += 1
-        return
 
     # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state"
     def _get_state(self):
diff --git a/apex/contrib/optimizers/fused_adam.py b/apex/contrib/optimizers/fused_adam.py
index 37a77b3cf..c0668737c 100644
--- a/apex/contrib/optimizers/fused_adam.py
+++ b/apex/contrib/optimizers/fused_adam.py
@@ -1,6 +1,8 @@
+import importlib
 import types
+
 import torch
-import importlib
+
 from apex.multi_tensor_apply import multi_tensor_applier
 
 
@@ -72,7 +74,7 @@ def __init__(
             weight_decay=weight_decay,
             max_grad_norm=max_grad_norm,
         )
-        super(FusedAdam, self).__init__(params, defaults)
+        super().__init__(params, defaults)
         self.eps_mode = 0 if eps_inside_sqrt else 1
 
     def step(self, closure=None, grads=None, output_params=None, scale=1.0, grad_norms=None):
@@ -104,18 +106,16 @@ def step(self, closure=None, grads=None, output_params=None, scale=1.0, grad_nor
             grads_group = [None] * len(self.param_groups)
         # backward compatibility
         # assuming a list/generator of parameter means single group
-        elif isinstance(grads, types.GeneratorType):
-            grads_group = [grads]
-        elif not isinstance(grads[0], list):
+        elif isinstance(grads, types.GeneratorType) or not isinstance(grads[0], list):
             grads_group = [grads]
         else:
             grads_group = grads
 
         if output_params is None:
             output_params_group = [None] * len(self.param_groups)
-        elif isinstance(output_params, types.GeneratorType):
-            output_params_group = [output_params]
-        elif not isinstance(output_params[0], list):
+        elif isinstance(output_params, types.GeneratorType) or not isinstance(
+            output_params[0], list
+        ):
             output_params_group = [output_params]
         else:
             output_params_group = output_params
diff --git a/apex/contrib/optimizers/fused_lamb.py b/apex/contrib/optimizers/fused_lamb.py
index d40cfeab7..2c9a0be64 100644
--- a/apex/contrib/optimizers/fused_lamb.py
+++ b/apex/contrib/optimizers/fused_lamb.py
@@ -1,6 +1,8 @@
-import torch
 import importlib
 import math
+
+import torch
+
 from apex.multi_tensor_apply import multi_tensor_applier
 
 
@@ -85,7 +87,7 @@ def __init__(
             grad_averaging=grad_averaging,
             max_grad_norm=max_grad_norm,
         )
-        super(FusedLAMB, self).__init__(params, defaults)
+        super().__init__(params, defaults)
         if multi_tensor_applier.available:
             import amp_C
 
@@ -105,7 +107,7 @@ def zero_grad(self):
                 for p in group["params"]:
                     p.grad = None
         else:
-            super(FusedLAMB, self).zero_grad()
+            super().zero_grad()
 
     def step(self, closure=None):
         """Performs a single optimization step.
diff --git a/apex/contrib/optimizers/fused_sgd.py b/apex/contrib/optimizers/fused_sgd.py
index e2acfcbaa..362460826 100644
--- a/apex/contrib/optimizers/fused_sgd.py
+++ b/apex/contrib/optimizers/fused_sgd.py
@@ -1,4 +1,5 @@
 import types
+
 import torch
 from torch.optim.optimizer import Optimizer, required
 
@@ -76,11 +77,11 @@ def __init__(
         materialize_master_grads=True,
     ):
         if lr is not required and lr < 0.0:
-            raise ValueError("Invalid learning rate: {}".format(lr))
+            raise ValueError(f"Invalid learning rate: {lr}")
         if momentum < 0.0:
-            raise ValueError("Invalid momentum value: {}".format(momentum))
+            raise ValueError(f"Invalid momentum value: {momentum}")
         if weight_decay < 0.0:
-            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
 
         defaults = dict(
             lr=lr,
@@ -91,7 +92,7 @@ def __init__(
         )
         if nesterov and (momentum <= 0 or dampening != 0):
             raise ValueError("Nesterov momentum requires a momentum and zero dampening")
-        super(FusedSGD, self).__init__(params, defaults)
+        super().__init__(params, defaults)
 
         self.wd_after_momentum = wd_after_momentum
 
@@ -105,7 +106,7 @@ def __init__(
             raise RuntimeError("apex.contrib.optimizers.FusedSGD requires cuda extensions")
 
     def __setstate__(self, state):
-        super(FusedSGD, self).__setstate__(state)
+        super().__setstate__(state)
         for group in self.param_groups:
             group.setdefault("nesterov", False)
 
@@ -155,9 +156,7 @@ def step(self, closure=None, grads=None, output_params=None, scale=1.0, grad_nor
             )
         # backward compatibility
         # assuming a list/generator of parameter means single group
-        elif isinstance(grads, types.GeneratorType):
-            grads_group = [grads]
-        elif not isinstance(grads[0], list):
+        elif isinstance(grads, types.GeneratorType) or not isinstance(grads[0], list):
             grads_group = [grads]
         else:
             grads_group = grads
@@ -168,9 +167,9 @@ def step(self, closure=None, grads=None, output_params=None, scale=1.0, grad_nor
                                with apex.contrib.optimizers.FP16_Optimizer \
                                which provides output_params."
             )
-        elif isinstance(output_params, types.GeneratorType):
-            output_params_group = [output_params]
-        elif not isinstance(output_params[0], list):
+        elif isinstance(output_params, types.GeneratorType) or not isinstance(
+            output_params[0], list
+        ):
             output_params_group = [output_params]
         else:
             output_params_group = output_params
diff --git a/apex/contrib/peer_memory/__init__.py b/apex/contrib/peer_memory/__init__.py
index 367dc5854..6aad74863 100644
--- a/apex/contrib/peer_memory/__init__.py
+++ b/apex/contrib/peer_memory/__init__.py
@@ -1,2 +1,2 @@
-from .peer_memory import PeerMemoryPool
 from .peer_halo_exchanger_1d import PeerHaloExchanger1d
+from .peer_memory import PeerMemoryPool
diff --git a/apex/contrib/peer_memory/peer_halo_exchanger_1d.py b/apex/contrib/peer_memory/peer_halo_exchanger_1d.py
index 8995e806d..e58deaf4e 100644
--- a/apex/contrib/peer_memory/peer_halo_exchanger_1d.py
+++ b/apex/contrib/peer_memory/peer_halo_exchanger_1d.py
@@ -1,5 +1,5 @@
-import torch
 import peer_memory_cuda as pm
+import torch
 
 
 class PeerHaloExchanger1d:
diff --git a/apex/contrib/peer_memory/peer_memory.py b/apex/contrib/peer_memory/peer_memory.py
index 72b1a1098..3d259b3ef 100644
--- a/apex/contrib/peer_memory/peer_memory.py
+++ b/apex/contrib/peer_memory/peer_memory.py
@@ -1,9 +1,9 @@
-import torch
 import numpy as np
 import peer_memory_cuda as pm
+import torch
 
 
-class PeerMemoryPool(object):
+class PeerMemoryPool:
     def __init__(self, static_size, dynamic_size, peer_ranks=None):
         rank = torch.distributed.get_rank()
         world_size = torch.distributed.get_world_size()
diff --git a/apex/contrib/sparsity/__init__.py b/apex/contrib/sparsity/__init__.py
index 661fd4ae9..4ec4ffca7 100644
--- a/apex/contrib/sparsity/__init__.py
+++ b/apex/contrib/sparsity/__init__.py
@@ -1,2 +1,2 @@
-from .sparse_masklib import create_mask
 from .asp import ASP
+from .sparse_masklib import create_mask
diff --git a/apex/contrib/sparsity/asp.py b/apex/contrib/sparsity/asp.py
index 7705a3f0b..38731e0df 100644
--- a/apex/contrib/sparsity/asp.py
+++ b/apex/contrib/sparsity/asp.py
@@ -1,7 +1,9 @@
 import types
+
 import torch
-from .sparse_masklib import create_mask
+
 from .permutation_lib import Permutation
+from .sparse_masklib import create_mask
 
 torchvision_imported = True
 try:
@@ -352,9 +354,7 @@ def compute_sparse_masks(cls):
                         time.perf_counter() - start_time_permute
                     )
                     print(
-                        "[compute_sparse_masks] Take {:.4f} seconds to find and apply permutations.".format(
-                            duration_build_offline_permutation_graph
-                        )
+                        f"[compute_sparse_masks] Take {duration_build_offline_permutation_graph:.4f} seconds to find and apply permutations."
                     )
 
             for module_name, module, p_name, p, mask, pruned in cls.__sparse_parameters:
@@ -451,20 +451,14 @@ def set_permutation_saving_params(
         print("\n[ASP][set_permutation_saving_param] Set permutation saving related parameters")
         print("\n[set_permutation_saving_param] Set permutation saving related parameters")
         cls.__allow_permutation = allow_permutation
-        print(
-            "[set_permutation_saving_param]\t Allow permutation: {}".format(cls.__allow_permutation)
-        )
+        print(f"[set_permutation_saving_param]\t Allow permutation: {cls.__allow_permutation}")
         cls.__save_permutation_graph = save_permutation_graph
         print(
-            "[set_permutation_saving_param]\t Save permutation graphs: {}".format(
-                cls.__save_permutation_graph
-            )
+            f"[set_permutation_saving_param]\t Save permutation graphs: {cls.__save_permutation_graph}"
         )
         cls.__permutation_output_dir = permutation_output_dir
         print(
-            "[set_permutation_saving_param]\t Permutation graphs saving dir: {}".format(
-                cls.__permutation_output_dir
-            )
+            f"[set_permutation_saving_param]\t Permutation graphs saving dir: {cls.__permutation_output_dir}"
         )
 
         Permutation.set_permutation_saving_params(
diff --git a/apex/contrib/sparsity/permutation_lib.py b/apex/contrib/sparsity/permutation_lib.py
index 375143410..b30bf7910 100644
--- a/apex/contrib/sparsity/permutation_lib.py
+++ b/apex/contrib/sparsity/permutation_lib.py
@@ -1,11 +1,12 @@
-import os
-import torch
+import builtins as __builtin__
+import io
 import json
+import os
 import string
 import time
+
 import numpy as np
-import builtins as __builtin__
-import io
+import torch
 
 try:
     from .permutation_search_kernels import (
@@ -160,9 +161,7 @@ def set_identical_seed(cls, identical_seed=1):
 
         if cls.__verbosity > 0:
             print(
-                "[set_identical_seed] Set the identical seed: {:} for all GPUs to make sure the same results generated in permutation search".format(
-                    identical_seed
-                )
+                f"[set_identical_seed] Set the identical seed: {identical_seed} for all GPUs to make sure the same results generated in permutation search"
             )
 
         cls.__manual_seed = identical_seed
@@ -302,9 +301,7 @@ def permute_model(
             )
             if cls.__verbosity > 0:
                 print(
-                    "\n[permute_model] Take {:.4f} seconds to finish search_for_good_permutation function.".format(
-                        duration_search_for_good_permutation
-                    )
+                    f"\n[permute_model] Take {duration_search_for_good_permutation:.4f} seconds to finish search_for_good_permutation function."
                 )
 
             fx_graph_after_sync_permutations = cls.sync_permutations(
@@ -342,11 +339,7 @@ def apply_permutation_in_C_dim(cls, node_name, permutation_sequence, dryrun):
         """This function is used to permutation for a node in C dim. (Only need to handle the weight of the node)"""
 
         if cls.__verbosity > 1 and dryrun:
-            print(
-                "[apply_permutation_in_C_dim] Permutation for node: '{:}' in C dim".format(
-                    node_name
-                )
-            )
+            print(f"[apply_permutation_in_C_dim] Permutation for node: '{node_name}' in C dim")
 
         if len(permutation_sequence) == 0:
             if cls.__verbosity >= 0:
@@ -361,9 +354,7 @@ def apply_permutation_in_C_dim(cls, node_name, permutation_sequence, dryrun):
             if node_name_matches(node_name, module_name):
                 if cls.__verbosity > 2 and dryrun:
                     print(
-                        "[apply_permutation_in_C_dim] find the node: '{:}' '{:}' in cls.__sparse_parameters, succeed to apply permutation in C dim.".format(
-                            node_name, p_name
-                        )
+                        f"[apply_permutation_in_C_dim] find the node: '{node_name}' '{p_name}' in cls.__sparse_parameters, succeed to apply permutation in C dim."
                     )
                 is_node_in_sparse_parameters = True
                 permutation_to_apply = permutation_sequence
@@ -397,9 +388,7 @@ def apply_permutation_in_C_dim(cls, node_name, permutation_sequence, dryrun):
                     ):
                         if cls.__verbosity > 3 and dryrun:
                             print(
-                                "[apply_permutation_in_C_dim] cannot find the node: '{:}' '{:}' in cls.__sparse_parameters, but can find in cls.__all_parameters.".format(
-                                    node_name, p_name_from_all_parameters
-                                )
+                                f"[apply_permutation_in_C_dim] cannot find the node: '{node_name}' '{p_name_from_all_parameters}' in cls.__sparse_parameters, but can find in cls.__all_parameters."
                             )
                         permutation_to_apply = permutation_sequence
                         if p_from_all_parameters.shape[1] != len(
@@ -424,17 +413,13 @@ def apply_permutation_in_C_dim(cls, node_name, permutation_sequence, dryrun):
                         success_permutation = True
                         if cls.__verbosity > 2 and dryrun:
                             print(
-                                "[apply_permutation_in_C_dim] cannot find the node: '{:}' in cls.__sparse_parameters, after trying with cls.__all_parameters, succeed to apply permutation in C dim.".format(
-                                    node_name
-                                )
+                                f"[apply_permutation_in_C_dim] cannot find the node: '{node_name}' in cls.__sparse_parameters, after trying with cls.__all_parameters, succeed to apply permutation in C dim."
                             )
             except:
                 success_permutation = False
                 if cls.__verbosity >= 0:
                     print(
-                        "ERROR: [apply_permutation_in_C_dim] cannot find the node: '{:}' in cls.__sparse_parameters, after trying with cls.__all_parameters, still fail to apply permutation in C dim.".format(
-                            node_name
-                        )
+                        f"ERROR: [apply_permutation_in_C_dim] cannot find the node: '{node_name}' in cls.__sparse_parameters, after trying with cls.__all_parameters, still fail to apply permutation in C dim."
                     )
         return success_permutation
 
@@ -483,11 +468,7 @@ def apply_permutation_in_K_dim(cls, node_name, permutation_sequence, fx_graph, d
         """This function is used to permutation for a node in K dim. (Need to handle the weight/bias/running_mean/running_var of the node)"""
 
         if cls.__verbosity > 1:
-            print(
-                "[apply_permutation_in_K_dim] Permutation for node: '{:}' in K dim".format(
-                    node_name
-                )
-            )
+            print(f"[apply_permutation_in_K_dim] Permutation for node: '{node_name}' in K dim")
 
         if len(permutation_sequence) == 0:
             if cls.__verbosity >= 0:
@@ -508,9 +489,7 @@ def apply_permutation_in_K_dim(cls, node_name, permutation_sequence, fx_graph, d
             if node_name_matches(node_name, module_name):
                 if cls.__verbosity > 1 and dryrun:
                     print(
-                        "[apply_permutation_in_K_dim] find the node: '{:}' with '{:}' in cls.__all_parameters, may succeed to apply permutation in K dim.".format(
-                            node_name, p_name
-                        )
+                        f"[apply_permutation_in_K_dim] find the node: '{node_name}' with '{p_name}' in cls.__all_parameters, may succeed to apply permutation in K dim."
                     )
                 is_node_in_all_parameters = True
                 permutation_to_apply = permutation_sequence
@@ -530,19 +509,12 @@ def apply_permutation_in_K_dim(cls, node_name, permutation_sequence, fx_graph, d
 
                     if cls.__verbosity > 1 and dryrun:
                         print(
-                            "[apply_permutation_in_K_dim] the node: '{:}' with shape: '{:}' required replicating the permutation sequence with len '{:}' {:} times to succeed in applying the permutation in the K dimension.".format(
-                                node_name,
-                                p.shape,
-                                len(permutation_sequence),
-                                p.shape[0] // len(permutation_sequence),
-                            )
+                            f"[apply_permutation_in_K_dim] the node: '{node_name}' with shape: '{p.shape}' required replicating the permutation sequence with len '{len(permutation_sequence)}' {p.shape[0] // len(permutation_sequence)} times to succeed in applying the permutation in the K dimension."
                         )
                 else:
                     if cls.__verbosity > 1 and dryrun:
                         print(
-                            "[apply_permutation_in_K_dim] the node: '{:}' with shape: '{:}', can match the size of permutation sequence with len: '{:}', succeed to apply permutation in K dim.".format(
-                                node_name, p.shape, len(permutation_sequence)
-                            )
+                            f"[apply_permutation_in_K_dim] the node: '{node_name}' with shape: '{p.shape}', can match the size of permutation sequence with len: '{len(permutation_sequence)}', succeed to apply permutation in K dim."
                         )
 
                 if not dryrun:
@@ -554,9 +526,7 @@ def apply_permutation_in_K_dim(cls, node_name, permutation_sequence, fx_graph, d
         if not is_node_in_all_parameters:
             if cls.__verbosity >= 0:
                 print(
-                    "ERROR: [apply_permutation_in _K_dim] cannot find the node: '{:}' in cls.__all_parameters, fail to apply permutation in K dim.".format(
-                        node_name
-                    )
+                    f"ERROR: [apply_permutation_in _K_dim] cannot find the node: '{node_name}' in cls.__all_parameters, fail to apply permutation in K dim."
                 )
             success_permutation = False
 
@@ -678,9 +648,7 @@ def find_permutation_for_matrix_group(cls, matrix_group):
 
         if cls.__verbosity > 1:
             print(
-                "\n[search_for_good_permutation] Original element abs sum: {:}, Pruned element abs sum: {:}, Diff ratio: {:}".format(
-                    original_magnitude, pruned_magnitude, diff_ratio
-                )
+                f"\n[search_for_good_permutation] Original element abs sum: {original_magnitude}, Pruned element abs sum: {pruned_magnitude}, Diff ratio: {diff_ratio}"
             )
 
         start_time_accelerated_search_for_good_permutation = time.perf_counter()
@@ -733,9 +701,7 @@ def find_permutation_for_matrix_group(cls, matrix_group):
                 matrix_group.cpu().detach().numpy()[:, group_permutation]
             )
             print(
-                "[search_for_good_permutation] Take {:.4f} seconds to finish accelerated_search_for_good_permutation function and with final magnitude {:}.".format(
-                    duration_accelerated_search_for_good_permutation, permuted_magnitude
-                )
+                f"[search_for_good_permutation] Take {duration_accelerated_search_for_good_permutation:.4f} seconds to finish accelerated_search_for_good_permutation function and with final magnitude {permuted_magnitude}."
             )
 
         return group_permutation, permutation_found
@@ -785,16 +751,12 @@ def collect_sparse_weights(cls, fx_graph, sibling_group, sibling_group_C_param):
                     except:
                         if cls.__verbosity >= 0:
                             print(
-                                "ERROR: [search_for_good_permutation][warning] cannot merge the weight for node: '{:}', with its weight shape: '{:}', the matrix_group shape: '{:}'.".format(
-                                    sibling, node_weight.size(), matrix_group.size()
-                                )
+                                f"ERROR: [search_for_good_permutation][warning] cannot merge the weight for node: '{sibling}', with its weight shape: '{node_weight.size()}', the matrix_group shape: '{matrix_group.size()}'."
                             )
                         continue
                 if cls.__verbosity > 2:
                     print(
-                        "[search_for_good_permutation] have merged the weight for node: '{:}', with its weight shape: '{:}', the matrix_group shape: '{:}'.".format(
-                            sibling, node_weight.size(), matrix_group.size()
-                        )
+                        f"[search_for_good_permutation] have merged the weight for node: '{sibling}', with its weight shape: '{node_weight.size()}', the matrix_group shape: '{matrix_group.size()}'."
                     )
             else:
                 if cls.__verbosity > 2:
@@ -1695,11 +1657,7 @@ def find_real_children(cls, fx_graph):
             node_children = fx_graph.get(node_name).get("children")
 
             if cls.__verbosity > 2:
-                print(
-                    "[find_real_children] node_name: '{:}', children: {:}".format(
-                        node_name, node_children
-                    )
-                )
+                print(f"[find_real_children] node_name: '{node_name}', children: {node_children}")
 
             real_children = cls.find_node_real_children(fx_graph, node_name, set())
 
@@ -1789,12 +1747,7 @@ def build_fx_graph(
             torch_version_minimum = torch_version.split(".")[2]
         if cls.__verbosity > 2:
             print(
-                "[build_fx_graph] The torch version is: {}, version major is: {}, version minor is: {}, version minimum is: {}".format(
-                    torch_version,
-                    torch_version_major,
-                    torch_version_minor,
-                    torch_version_minimum,
-                )
+                f"[build_fx_graph] The torch version is: {torch_version}, version major is: {torch_version_major}, version minor is: {torch_version_minor}, version minimum is: {torch_version_minimum}"
             )
 
         if torch_version_major >= 2 or (torch_version_major >= 1 and torch_version_minor >= 8):
@@ -1828,7 +1781,7 @@ def build_fx_graph(
         module_name_K_dict = {}
         for name, mod in model.named_modules():
             if cls.__verbosity > 1:
-                print("[build_fx_graph] module_name: {}, module type: {}".format(name, type(mod)))
+                print(f"[build_fx_graph] module_name: {name}, module type: {type(mod)}")
             module_name_type_dict[name] = str(type(mod)).split("'")[1]
             try:
                 module_name_C_dict[name] = str(mod.in_channels)
@@ -1856,9 +1809,7 @@ def build_fx_graph(
                 module_name_group_conv_dict[name] = str(mod.groups)
                 if cls.__verbosity > 1:
                     print(
-                        "[build_fx_graph] this module has 'group' param with value: {}".format(
-                            mod.groups
-                        )
+                        f"[build_fx_graph] this module has 'group' param with value: {mod.groups}"
                     )
             except:
                 module_name_group_conv_dict[name] = "None"
@@ -1871,11 +1822,11 @@ def build_fx_graph(
         for node in graph_module.graph.nodes:
             if node.op == "placeholder":
                 if cls.__verbosity > 2:
-                    print("[build_fx_graph] This is the 'input' node: {:}".format(node.target))
+                    print(f"[build_fx_graph] This is the 'input' node: {node.target}")
                 continue
             elif node.op == "get_attr":
                 if cls.__verbosity > 2:
-                    print("[build_fx_graph] This is the 'get_attr' node: {:}".format(node.target))
+                    print(f"[build_fx_graph] This is the 'get_attr' node: {node.target}")
                 node_parent, node_children = get_node_parent_children(node)
                 converted_node_name = convert_fx_node_name(node.target)
 
@@ -1909,9 +1860,7 @@ def fetch_attr(target: str, mod):
                 converted_node_name = convert_fx_node_name(node.name)
                 if cls.__verbosity > 2:
                     print(
-                        "[build_fx_graph] This is the 'call_function' node: {:}, its parent list: {:}, its children list: {:}".format(
-                            converted_node_name, node_parent, node_children
-                        )
+                        f"[build_fx_graph] This is the 'call_function' node: {converted_node_name}, its parent list: {node_parent}, its children list: {node_children}"
                     )
                 network_fx_graph[converted_node_name] = {}
                 network_fx_graph[converted_node_name]["parents"] = node_parent
@@ -1938,9 +1887,7 @@ def fetch_attr(target: str, mod):
                 converted_node_name = convert_fx_node_name(node.name)
                 if cls.__verbosity > 2:
                     print(
-                        "[build_fx_graph] This is the 'call_method' node: {:}, its parent list: {:}, its children list: {:}".format(
-                            converted_node_name, node_parent, node_children
-                        )
+                        f"[build_fx_graph] This is the 'call_method' node: {converted_node_name}, its parent list: {node_parent}, its children list: {node_children}"
                     )
                 network_fx_graph[converted_node_name] = {}
                 network_fx_graph[converted_node_name]["parents"] = node_parent
@@ -1955,18 +1902,14 @@ def fetch_attr(target: str, mod):
                 if converted_node_name != node.target:
                     if cls.__verbosity > 2:
                         print(
-                            "[build_fx_graph][warning] The target name from Torch.FX is '{:}', the manually converted node name is '{:}', not the same one, choose the converted node name".format(
-                                node.target, converted_node_name
-                            )
+                            f"[build_fx_graph][warning] The target name from Torch.FX is '{node.target}', the manually converted node name is '{converted_node_name}', not the same one, choose the converted node name"
                         )
 
                 # assume the modules share the same target name have the same type, because converted_node_name may not be obtained by model.named_modules(), like some ReLU (defined in forward function)
                 node_type = module_name_type_dict[node.target]
                 if cls.__verbosity > 2:
                     print(
-                        "[build_fx_graph] This is the 'call_module' node: {:}, its parent list: {:}, its children list: {:}, its type: {:}".format(
-                            converted_node_name, node_parent, node_children, node_type
-                        )
+                        f"[build_fx_graph] This is the 'call_module' node: {converted_node_name}, its parent list: {node_parent}, its children list: {node_children}, its type: {node_type}"
                     )
                 network_fx_graph[converted_node_name] = {}
                 network_fx_graph[converted_node_name]["parents"] = node_parent
@@ -1981,7 +1924,7 @@ def fetch_attr(target: str, mod):
 
             elif node.op == "output":
                 if cls.__verbosity > 2:
-                    print("[build_fx_graph] This is the 'output' node: {:}".format(node.target))
+                    print(f"[build_fx_graph] This is the 'output' node: {node.target}")
                 continue
 
         if dump_fx_graph:
@@ -1999,9 +1942,10 @@ def fetch_attr(target: str, mod):
     def trace_and_print_raw_fx_graph(cls, model, print_tabular=False, generate_python_code=False):
         """This function is used to find and print the intermediate representation (IR) - Graph representation with Torch.FX features."""
 
-        from torch.fx import symbolic_trace
         import traceback
 
+        from torch.fx import symbolic_trace
+
         # Symbolic tracing frontend - captures the semantics of the module
         try:
             symbolic_traced: torch.fx.GraphModule = symbolic_trace(model)
diff --git a/apex/contrib/sparsity/permutation_search_kernels/call_permutation_search_kernels.py b/apex/contrib/sparsity/permutation_search_kernels/call_permutation_search_kernels.py
index 774327193..b5c0178ab 100644
--- a/apex/contrib/sparsity/permutation_search_kernels/call_permutation_search_kernels.py
+++ b/apex/contrib/sparsity/permutation_search_kernels/call_permutation_search_kernels.py
@@ -1,6 +1,7 @@
 import numpy as np
-from .permutation_utilities import *
+
 from .exhaustive_search import Exhaustive_Search
+from .permutation_utilities import *
 
 
 def accelerated_search_for_good_permutation(matrix_group, options=None, verbosity=0):
@@ -11,9 +12,7 @@ def accelerated_search_for_good_permutation(matrix_group, options=None, verbosit
     input_matrix = matrix_group.cpu().detach().numpy()
     if verbosity > 1:
         print(
-            "\n[accelerated_search_for_good_permutation] input matrix shape: '{:}'.".format(
-                input_matrix.shape
-            )
+            f"\n[accelerated_search_for_good_permutation] input matrix shape: '{input_matrix.shape}'."
         )
 
     result = np.copy(input_matrix)
@@ -78,9 +77,7 @@ def accelerated_search_for_good_permutation(matrix_group, options=None, verbosit
         duration = time.perf_counter() - start_time
         if verbosity > 1:
             print(
-                "\tFinally swap {} channel pairs until the search time limit expires.".format(
-                    real_swap_num
-                )
+                f"\tFinally swap {real_swap_num} channel pairs until the search time limit expires."
             )
     elif (
         options["strategy"] == "user defined"
@@ -97,9 +94,7 @@ def accelerated_search_for_good_permutation(matrix_group, options=None, verbosit
 
     if verbosity > 1:
         print(
-            "[accelerated_search_for_good_permutation] Take {:.4f} seconds to search the permutation sequence.".format(
-                duration
-            )
+            f"[accelerated_search_for_good_permutation] Take {duration:.4f} seconds to search the permutation sequence."
         )
 
     return permutation_sequence
diff --git a/apex/contrib/sparsity/permutation_search_kernels/exhaustive_search.py b/apex/contrib/sparsity/permutation_search_kernels/exhaustive_search.py
index 1fc168650..b84d2ef4e 100644
--- a/apex/contrib/sparsity/permutation_search_kernels/exhaustive_search.py
+++ b/apex/contrib/sparsity/permutation_search_kernels/exhaustive_search.py
@@ -102,7 +102,7 @@ def generate_all_unique_combinations(C, M, must_use_all_groups=False):
 def predict_unique_combinations(C, M):
     assert C % M == 0
     G = int(C / M)
-    return int(int(math.factorial(C)) / (int(math.pow(math.factorial(M), G)) * math.factorial(G)))
+    return int(math.factorial(C) / (int(math.pow(math.factorial(M), G)) * math.factorial(G)))
 
 
 #################################################################
@@ -350,7 +350,7 @@ def use_stripe_map(matrix, group_width, stripe_map, stripe_ids, perm_map, permut
                     break
             # if it's not, then it changed
             if changed:
-                used_stripes.append(stripe_group[s])
+                used_stripes.append(stripe)
 
             matrix[..., stripe * group_width : stripe * group_width + group_width] = sub_result[
                 ..., s * group_width : s * group_width + group_width
diff --git a/apex/contrib/sparsity/permutation_search_kernels/permutation_utilities.py b/apex/contrib/sparsity/permutation_search_kernels/permutation_utilities.py
index 7d253fa0b..bad25bccd 100644
--- a/apex/contrib/sparsity/permutation_search_kernels/permutation_utilities.py
+++ b/apex/contrib/sparsity/permutation_search_kernels/permutation_utilities.py
@@ -1,6 +1,7 @@
-import numpy as np
-import subprocess
 import math
+import subprocess
+
+import numpy as np
 
 gpus_tested = False
 gpus_found = 0
@@ -526,7 +527,7 @@ def move_permutation_towards(B, A, debug=False):
         cur_entry = wrong_entries[we]
         # if debug:
         #    print(f"\tMPT: checking {cur_entry} for complement")
-        for we2 in range(0, len(wrong_entries)):
+        for we2 in range(len(wrong_entries)):
             pos_swap = wrong_entries[we2]
             # if debug:
             #    print(f"\t\tMPT: is {pos_swap}?")
diff --git a/apex/contrib/sparsity/permutation_tests/permutation_test.py b/apex/contrib/sparsity/permutation_tests/permutation_test.py
index c8071c738..80bd203fb 100644
--- a/apex/contrib/sparsity/permutation_tests/permutation_test.py
+++ b/apex/contrib/sparsity/permutation_tests/permutation_test.py
@@ -1,16 +1,17 @@
-import numpy as np
-import time
 import sys
+import time
+
+import numpy as np
 
 # permutation-specifics
 sys.path.append("../")
-from permutation_search_kernels.permutation_utilities import *
-from permutation_search_kernels.exhaustive_search import Exhaustive_Search
-from permutation_search_kernels.channel_swap import Channel_Swap
-
 # Arguments
 import argparse
 
+from permutation_search_kernels.channel_swap import Channel_Swap
+from permutation_search_kernels.exhaustive_search import Exhaustive_Search
+from permutation_search_kernels.permutation_utilities import *
+
 
 def str2bool(v):
     if isinstance(v, bool):
diff --git a/apex/contrib/sparsity/sparse_masklib.py b/apex/contrib/sparsity/sparse_masklib.py
index d36af16a7..81ce33fed 100644
--- a/apex/contrib/sparsity/sparse_masklib.py
+++ b/apex/contrib/sparsity/sparse_masklib.py
@@ -1,9 +1,9 @@
-import sys
-import torch
-import numpy as np
 import collections
+import sys
 from itertools import permutations
 
+import numpy as np
+import torch
 
 """ compute density (helper fn to compute % NNZs in a tensor) """
 
diff --git a/apex/contrib/sparsity/test/checkpointing_test_part1.py b/apex/contrib/sparsity/test/checkpointing_test_part1.py
index 193c72ce2..58e3f8185 100644
--- a/apex/contrib/sparsity/test/checkpointing_test_part1.py
+++ b/apex/contrib/sparsity/test/checkpointing_test_part1.py
@@ -1,8 +1,9 @@
 from collections import OrderedDict
 
 import torch
-from apex.optimizers import FusedAdam
+
 from apex.contrib.sparsity import ASP
+from apex.optimizers import FusedAdam
 
 
 def build_model(args):
diff --git a/apex/contrib/sparsity/test/checkpointing_test_part2.py b/apex/contrib/sparsity/test/checkpointing_test_part2.py
index eb77b6e88..14fb67f87 100644
--- a/apex/contrib/sparsity/test/checkpointing_test_part2.py
+++ b/apex/contrib/sparsity/test/checkpointing_test_part2.py
@@ -1,8 +1,9 @@
 from collections import OrderedDict
 
 import torch
-from apex.optimizers import FusedAdam
+
 from apex.contrib.sparsity import ASP
+from apex.optimizers import FusedAdam
 
 
 def build_model(args):
diff --git a/apex/contrib/sparsity/test/checkpointing_test_reference.py b/apex/contrib/sparsity/test/checkpointing_test_reference.py
index 21b50efb6..b4ea2116d 100644
--- a/apex/contrib/sparsity/test/checkpointing_test_reference.py
+++ b/apex/contrib/sparsity/test/checkpointing_test_reference.py
@@ -1,8 +1,9 @@
 from collections import OrderedDict
 
 import torch
-from apex.optimizers import FusedAdam
+
 from apex.contrib.sparsity import ASP
+from apex.optimizers import FusedAdam
 
 #
 # Reference run for checkpointing test (part1 + part2)
diff --git a/apex/contrib/sparsity/test/test_permutation_application.py b/apex/contrib/sparsity/test/test_permutation_application.py
index 411c93bc6..0c3563f7d 100644
--- a/apex/contrib/sparsity/test/test_permutation_application.py
+++ b/apex/contrib/sparsity/test/test_permutation_application.py
@@ -1,5 +1,6 @@
 import torch
 import torch.onnx
+
 from apex.contrib.sparsity.permutation_lib import Permutation
 
 """
@@ -821,7 +822,7 @@ def test_model(model, tag, verbosity=0, save_onnx=False):
                 ),
             ):
                 allowed_names = ("weight",)
-                if type(module) in module_to_params.keys():
+                if type(module) in module_to_params:
                     allowed_names = module_to_params[type(module)]
 
                 if p_name not in allowed_names:
diff --git a/apex/contrib/sparsity/test/toy_problem.py b/apex/contrib/sparsity/test/toy_problem.py
index 9f766a330..7761a129a 100644
--- a/apex/contrib/sparsity/test/toy_problem.py
+++ b/apex/contrib/sparsity/test/toy_problem.py
@@ -1,8 +1,9 @@
 from collections import OrderedDict
 
 import torch
-from apex.optimizers import FusedAdam
+
 from apex.contrib.sparsity import ASP
+from apex.optimizers import FusedAdam
 
 
 def build_model(args):
diff --git a/apex/contrib/test/bottleneck/test_bottleneck_module.py b/apex/contrib/test/bottleneck/test_bottleneck_module.py
index 634f1e741..a76f748d6 100644
--- a/apex/contrib/test/bottleneck/test_bottleneck_module.py
+++ b/apex/contrib/test/bottleneck/test_bottleneck_module.py
@@ -7,8 +7,7 @@
 
 SKIP_TEST = None
 try:
-    from apex.contrib.bottleneck import Bottleneck, SpatialBottleneck
-    from apex.contrib.bottleneck import HaloExchangerPeer
+    from apex.contrib.bottleneck import Bottleneck, HaloExchangerPeer, SpatialBottleneck
     from apex.contrib.peer_memory import PeerMemoryPool
 except ImportError as e:
     SKIP_TEST = e
diff --git a/apex/contrib/test/conv_bias_relu/test_conv_bias_relu.py b/apex/contrib/test/conv_bias_relu/test_conv_bias_relu.py
index 9fb249651..1102800b4 100644
--- a/apex/contrib/test/conv_bias_relu/test_conv_bias_relu.py
+++ b/apex/contrib/test/conv_bias_relu/test_conv_bias_relu.py
@@ -9,9 +9,9 @@
 HAS_CONV_BIAS_RELU = None
 try:
     from apex.contrib.conv_bias_relu import (
-        ConvBiasReLU,
         ConvBias,
         ConvBiasMaskReLU,
+        ConvBiasReLU,
         ConvFrozenScaleBiasReLU,
     )
 except ImportError:
@@ -100,20 +100,9 @@ def setUp(self, seed=0):
         self.conv2_ = copy.deepcopy(self.conv2)
 
         print()
+        print(f"> input=[{self.batch_size}, {self.in_channels}, {self.in_height}, {self.in_width}]")
         print(
-            "> input=[{}, {}, {}, {}]".format(
-                self.batch_size, self.in_channels, self.in_height, self.in_width
-            )
-        )
-        print(
-            "> kernel=[{}, {}, {}, {}], stride={}, pad={}".format(
-                self.out_channels,
-                self.in_channels,
-                self.conv_kernel_size,
-                self.conv_kernel_size,
-                self.conv_stride,
-                self.conv_pad,
-            )
+            f"> kernel=[{self.out_channels}, {self.in_channels}, {self.conv_kernel_size}, {self.conv_kernel_size}], stride={self.conv_stride}, pad={self.conv_pad}"
         )
 
     def test_conv_bias_relu(self):
diff --git a/apex/contrib/test/cudnn_gbn/test_cudnn_gbn_with_two_gpus.py b/apex/contrib/test/cudnn_gbn/test_cudnn_gbn_with_two_gpus.py
index bbef3cd65..b9929868a 100644
--- a/apex/contrib/test/cudnn_gbn/test_cudnn_gbn_with_two_gpus.py
+++ b/apex/contrib/test/cudnn_gbn/test_cudnn_gbn_with_two_gpus.py
@@ -3,7 +3,7 @@
 import unittest
 
 import torch
-import torch.nn as nn
+from torch import nn
 from torch.testing._internal import common_utils
 
 SKIP_TEST = None
@@ -103,7 +103,7 @@ def world_size(self) -> int:
     def _test_cudnn_gbn(
         self,
         num_layers: int,
-        shape: typing.List[int],
+        shape: list[int],
         *,
         memory_format: torch.memory_format = torch.channels_last,
     ) -> None:
diff --git a/apex/contrib/test/fused_dense/test_fused_dense.py b/apex/contrib/test/fused_dense/test_fused_dense.py
index d25f969c1..33b1e2c8c 100644
--- a/apex/contrib/test/fused_dense/test_fused_dense.py
+++ b/apex/contrib/test/fused_dense/test_fused_dense.py
@@ -1,5 +1,5 @@
-import unittest
 import os
+import unittest
 
 import torch
 from torch.testing._internal import common_utils
diff --git a/apex/contrib/test/group_norm/test_group_norm.py b/apex/contrib/test/group_norm/test_group_norm.py
index 99bc7072b..3f4f1ca9d 100644
--- a/apex/contrib/test/group_norm/test_group_norm.py
+++ b/apex/contrib/test/group_norm/test_group_norm.py
@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# coding: utf-8
 
 #
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
@@ -10,16 +9,19 @@
 import importlib
 import pathlib
 import sys
-import torch
 import unittest
 
+import torch
+
 SKIP_TEST = None
 try:
-    from apex.contrib.group_norm.group_norm import cuda_group_norm_nhwc_one_pass
-    from apex.contrib.group_norm.group_norm import cuda_group_norm_nhwc_two_pass
-    from apex.contrib.group_norm.group_norm import cuda_group_norm_v2_nhwc
-    from apex.contrib.group_norm.group_norm import get_cc_and_sm_count
     from apex.contrib.group_norm import GroupNorm
+    from apex.contrib.group_norm.group_norm import (
+        cuda_group_norm_nhwc_one_pass,
+        cuda_group_norm_nhwc_two_pass,
+        cuda_group_norm_v2_nhwc,
+        get_cc_and_sm_count,
+    )
 except ImportError as e:
     SKIP_TEST = e
 
diff --git a/apex/contrib/test/layer_norm/test_fast_layer_norm.py b/apex/contrib/test/layer_norm/test_fast_layer_norm.py
index c85afa612..0d58fb91d 100644
--- a/apex/contrib/test/layer_norm/test_fast_layer_norm.py
+++ b/apex/contrib/test/layer_norm/test_fast_layer_norm.py
@@ -5,8 +5,9 @@
 
 SKIP_TEST = None
 try:
-    from apex.contrib.layer_norm.layer_norm import FastLayerNorm
     import fast_layer_norm as fln
+
+    from apex.contrib.layer_norm.layer_norm import FastLayerNorm
 except ImportError as e:
     SKIP_TEST = e
 
@@ -98,9 +99,7 @@ def benchmark_(S, B, hidden_size, itype, wtype, runs=100):
         ms_fwd = timer.millis() / runs
 
         print(
-            "[FWD] Time: {:.4f}ms Throughput: {:.4f} GB/sec".format(
-                ms_fwd, total_bytes_fwd * 1e-6 / ms_fwd
-            )
+            f"[FWD] Time: {ms_fwd:.4f}ms Throughput: {total_bytes_fwd * 1e-6 / ms_fwd:.4f} GB/sec"
         )
 
         timer.start()
@@ -132,9 +131,7 @@ def benchmark_(S, B, hidden_size, itype, wtype, runs=100):
         ms_bwd = timer.millis() / runs
 
         print(
-            "[BWD] Time: {:.4f}ms Throughput: {:.4f} GB/sec".format(
-                ms_bwd, total_bytes_bwd * 1e-6 / ms_bwd
-            )
+            f"[BWD] Time: {ms_bwd:.4f}ms Throughput: {total_bytes_bwd * 1e-6 / ms_bwd:.4f} GB/sec"
         )
 
 
diff --git a/apex/contrib/test/openfold_triton/test_fused_adam_swa.py b/apex/contrib/test/openfold_triton/test_fused_adam_swa.py
index 3f6dfa2e7..334a2ae00 100644
--- a/apex/contrib/test/openfold_triton/test_fused_adam_swa.py
+++ b/apex/contrib/test/openfold_triton/test_fused_adam_swa.py
@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
-from itertools import chain
 import random
 import unittest
+from itertools import chain
 
 import torch
-import torch.nn as nn
+from torch import nn
 
 SKIP_TEST = None
 try:
@@ -32,7 +32,7 @@ class AlphaFoldSWA(nn.Module):
     """AlphaFold SWA (Stochastic Weight Averaging) module wrapper."""
 
     def __init__(self, alphafold: nn.Module, enabled: bool, decay_rate: float) -> None:
-        super(AlphaFoldSWA, self).__init__()
+        super().__init__()
         if enabled:
             self.averaged_model = torch.optim.swa_utils.AveragedModel(
                 model=alphafold,
diff --git a/apex/contrib/test/openfold_triton/test_openfold_mha.py b/apex/contrib/test/openfold_triton/test_openfold_mha.py
index b1e6717be..c088136d2 100644
--- a/apex/contrib/test/openfold_triton/test_openfold_mha.py
+++ b/apex/contrib/test/openfold_triton/test_openfold_mha.py
@@ -1,8 +1,9 @@
 import math
 import random
+import unittest
 from typing import Optional
+
 import torch
-import unittest
 
 SKIP_TEST = None
 try:
@@ -16,7 +17,7 @@ def openfold_attention_eager(
     key: torch.Tensor,
     value: torch.Tensor,
     mask: torch.Tensor,
-    bias: Optional[torch.Tensor],
+    bias: torch.Tensor | None,
     inf: float,
 ) -> torch.Tensor:
     # query:  [*, num_heads, Q, c_hidden]
diff --git a/apex/contrib/test/openfold_triton/test_sync_triton_auto_tune_cache_across_gpus.py b/apex/contrib/test/openfold_triton/test_sync_triton_auto_tune_cache_across_gpus.py
index 0ec5aee35..d8cb70d29 100644
--- a/apex/contrib/test/openfold_triton/test_sync_triton_auto_tune_cache_across_gpus.py
+++ b/apex/contrib/test/openfold_triton/test_sync_triton_auto_tune_cache_across_gpus.py
@@ -2,17 +2,17 @@
 
 import torch
 import torch.distributed as dist
-from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.common_distributed import (
     MultiProcessTestCase,
     requires_nccl,
     skip_if_lt_x_gpu,
 )
+from torch.testing._internal.common_utils import run_tests
 
 from apex.contrib.openfold_triton import (
     LayerNormSmallShapeOptImpl,
-    sync_triton_auto_tune_cache_across_gpus,
     _tuneable_triton_kernels,
+    sync_triton_auto_tune_cache_across_gpus,
 )
 
 
diff --git a/apex/contrib/test/optimizers/test_dist_adam.py b/apex/contrib/test/optimizers/test_dist_adam.py
index 1a490b8dd..9b7995f70 100644
--- a/apex/contrib/test/optimizers/test_dist_adam.py
+++ b/apex/contrib/test/optimizers/test_dist_adam.py
@@ -1,9 +1,9 @@
-from contextlib import contextmanager
 import io
-from typing import Callable, Optional
 import unittest
 import warnings
-from contextlib import nullcontext
+from collections.abc import Callable
+from contextlib import contextmanager, nullcontext
+from typing import Optional
 
 import torch
 from torch.testing._internal import common_utils
@@ -37,11 +37,11 @@ def make_models(
     lr: float = 0.1,
     adam_w_mode: bool = True,
     model_dtype: torch.dtype = torch.float32,
-    optim_dtype: Optional[torch.dtype] = None,
-    grad_sync_dtype: Optional[torch.dtype] = None,
-    param_sync_dtype: Optional[torch.dtype] = None,
+    optim_dtype: torch.dtype | None = None,
+    grad_sync_dtype: torch.dtype | None = None,
+    param_sync_dtype: torch.dtype | None = None,
     device: torch.device = "cuda",
-    process_group: Optional[torch.distributed.ProcessGroup] = None,
+    process_group: torch.distributed.ProcessGroup | None = None,
     average_grad_sync: bool = True,
     overlap_communication: bool = True,
     bucket_cap_mb: float = 71 / (4 * 1024 * 1024),
@@ -121,8 +121,8 @@ class TestDistributedFusedAdam(NcclDistributedTestBase):
 
     def test_matches_pytorch(
         self,
-        rtol: Optional[float] = None,
-        atol: Optional[float] = None,
+        rtol: float | None = None,
+        atol: float | None = None,
         num_layers: int = 11,
         layer_size: int = 7,
         batch_size: int = 3,
@@ -132,9 +132,9 @@ def test_matches_pytorch(
         overlap_communication: bool = True,
         use_nosync: bool = True,
         model_dtype: torch.dtype = torch.float32,
-        optim_dtype: Optional[torch.dtype] = None,
-        grad_sync_dtype: Optional[torch.dtype] = None,
-        param_sync_dtype: Optional[torch.dtype] = None,
+        optim_dtype: torch.dtype | None = None,
+        grad_sync_dtype: torch.dtype | None = None,
+        param_sync_dtype: torch.dtype | None = None,
         device: torch.device = "cuda",
         bucket_cap_mb: float = 71 / (4 * 1024 * 1024),
         contiguous_buffers: bool = False,
@@ -142,7 +142,7 @@ def test_matches_pytorch(
         store_param_remainders: bool = False,
         with_scaled_states: bool = False,
         nccl_ub: bool = False,
-        init_optim_func: Optional[Callable[[DistributedFusedAdam], None]] = None,
+        init_optim_func: Callable[[DistributedFusedAdam], None] | None = None,
         with_cuda_graph: bool = False,
     ):
         torch.manual_seed(self.seed + self.rank)
@@ -491,15 +491,15 @@ def test_grad_scaler(self):
 
     def test_checkpoint(
         self,
-        rtol: Optional[float] = None,
-        atol: Optional[float] = None,
+        rtol: float | None = None,
+        atol: float | None = None,
         num_layers: int = 2,
         layer_size: int = 2,
         num_steps: int = 3,
-        save_group_size: Optional[int] = None,
-        load_group_size: Optional[int] = None,
-        save_model_kwargs: Optional[dict] = None,
-        load_model_kwargs: Optional[dict] = None,
+        save_group_size: int | None = None,
+        load_group_size: int | None = None,
+        save_model_kwargs: dict | None = None,
+        load_model_kwargs: dict | None = None,
     ):
         """Test state_dict and load_state_dict functions
 
@@ -595,8 +595,8 @@ def make_global_batch() -> torch.Tensor:
 
         def to_local_batch(
             global_batch: torch.Tensor,
-            group: Optional[torch.distributed.ProcessGroup],
-        ) -> Optional[torch.Tensor]:
+            group: torch.distributed.ProcessGroup | None,
+        ) -> torch.Tensor | None:
             """Get local portion of tensor that is replicated across all ranks"""
             group_size = torch.distributed.get_world_size(group)
             if group_size < 0:
@@ -608,7 +608,7 @@ def to_local_batch(
 
         def to_global_batch(
             local_batch: torch.Tensor,
-            group: Optional[torch.distributed.ProcessGroup],
+            group: torch.distributed.ProcessGroup | None,
         ) -> torch.Tensor:
             """Gather distributed tensor and broadcast to all ranks"""
 
diff --git a/apex/contrib/test/optimizers/test_distributed_fused_lamb.py b/apex/contrib/test/optimizers/test_distributed_fused_lamb.py
index 766311313..bfa3619c6 100644
--- a/apex/contrib/test/optimizers/test_distributed_fused_lamb.py
+++ b/apex/contrib/test/optimizers/test_distributed_fused_lamb.py
@@ -2,8 +2,8 @@
 
 import torch
 from torch.cuda.amp import GradScaler
-from torch.testing._internal import common_utils
 from torch.distributed.distributed_c10d import _coalescing_manager
+from torch.testing._internal import common_utils
 
 from apex.contrib.optimizers.distributed_fused_lamb import DistributedFusedLAMB
 from apex.distributed_testing.distributed_test_base import NcclDistributedTestBase
@@ -28,7 +28,7 @@ def init_weights(m):
 
 class ModelFoo(torch.nn.Module):
     def __init__(self):
-        super(ModelFoo, self).__init__()
+        super().__init__()
         self.linear = torch.nn.Linear(128, 128, bias=False)
         self.loss = torch.nn.MSELoss()
 
diff --git a/apex/contrib/test/peer_memory/test_peer_halo_exchange_module.py b/apex/contrib/test/peer_memory/test_peer_halo_exchange_module.py
index 6d85c0edb..bd1bcfed4 100644
--- a/apex/contrib/test/peer_memory/test_peer_halo_exchange_module.py
+++ b/apex/contrib/test/peer_memory/test_peer_halo_exchange_module.py
@@ -7,7 +7,7 @@
 from apex.distributed_testing.distributed_test_base import NcclDistributedTestBase
 
 try:
-    from apex.contrib.peer_memory import PeerMemoryPool, PeerHaloExchanger1d
+    from apex.contrib.peer_memory import PeerHaloExchanger1d, PeerMemoryPool
 except ImportError as e:
     SKIP_TEST = e
 
diff --git a/apex/contrib/test/xentropy/test_label_smoothing.py b/apex/contrib/test/xentropy/test_label_smoothing.py
index fe0cd7cdc..a168fccd0 100644
--- a/apex/contrib/test/xentropy/test_label_smoothing.py
+++ b/apex/contrib/test/xentropy/test_label_smoothing.py
@@ -1,9 +1,8 @@
-import unittest
 import random
 import time
+import unittest
 
 import numpy as np
-
 import torch
 
 SKIP_TEST = None
@@ -59,11 +58,7 @@ def print_max_diff_elem(self, ref, tst):
         ref, tst = ref.flatten(), tst.flatten()
         diff = (ref - tst).abs().max()
         idx = (ref - tst).abs().argmax()
-        print(
-            "Max atol idx: {}, diff: {:.6f}, ref: {:.6f}, tst: {:.6f}".format(
-                idx, diff, ref[idx], tst[idx]
-            )
-        )
+        print(f"Max atol idx: {idx}, diff: {diff:.6f}, ref: {ref[idx]:.6f}, tst: {tst[idx]:.6f}")
 
     def _test_label_smoothing_function(self, dtype):
         # Set label smoothing configuration
@@ -124,9 +119,7 @@ def test_label_smoothing_perf(self):
             loss.backward()
         torch.cuda.synchronize()
         print(
-            "Raw time {:.2f} s elapsed for {} iterations, norm {:.4f}".format(
-                time.time() - ts, iters, logits.grad.norm()
-            )
+            f"Raw time {time.time() - ts:.2f} s elapsed for {iters} iterations, norm {logits.grad.norm():.4f}"
         )
 
         # Run optimized softmax cross entropy with label smoothing
@@ -139,9 +132,7 @@ def test_label_smoothing_perf(self):
             loss.backward()
         torch.cuda.synchronize()
         print(
-            "Opt time {:.2f} s elapsed for {} iterations, norm {:.4f}".format(
-                time.time() - ts, iters, logits.grad.norm()
-            )
+            f"Opt time {time.time() - ts:.2f} s elapsed for {iters} iterations, norm {logits.grad.norm():.4f}"
         )
 
 
diff --git a/apex/contrib/torchsched/__init__.py b/apex/contrib/torchsched/__init__.py
index a91267a06..4f7a03cf1 100644
--- a/apex/contrib/torchsched/__init__.py
+++ b/apex/contrib/torchsched/__init__.py
@@ -6,8 +6,7 @@
 
 import torch
 import torch._inductor
-from torch._dynamo import list_backends
-from torch._dynamo import register_backend
+from torch._dynamo import list_backends, register_backend
 from torch._inductor.compile_fx import compile_fx_inner
 
 from .backend import get_backend
diff --git a/apex/contrib/torchsched/backend.py b/apex/contrib/torchsched/backend.py
index 3a14d3d90..437f498a2 100644
--- a/apex/contrib/torchsched/backend.py
+++ b/apex/contrib/torchsched/backend.py
@@ -4,23 +4,19 @@
 
 import functools
 from copy import copy
-from typing import TYPE_CHECKING
-from typing import ParamSpec
-from typing import TypeVar
+from typing import TYPE_CHECKING, ParamSpec, TypeVar
 
 if TYPE_CHECKING:
     from collections.abc import Callable
     from types import NotImplementedType
 
 import torch
-from torch import Tensor
-from torch import _TorchCompileInductorWrapper
+from torch import Tensor, _TorchCompileInductorWrapper
 from torch._dynamo import lookup_backend
-from torch._inductor.compile_fx import compile_fx
-from torch._inductor.compile_fx import compile_fx_inner
+from torch._inductor.compile_fx import compile_fx, compile_fx_inner
 from torch._inductor.decomposition import select_decomp_table
 
-import apex.contrib.torchsched.config as config
+from apex.contrib.torchsched import config
 from apex.contrib.torchsched.inductor import patch_graph_lowering
 from apex.contrib.torchsched.passes import pre_grad_custom_pass
 
diff --git a/apex/contrib/torchsched/config.py b/apex/contrib/torchsched/config.py
index 7b9befd9d..42034d492 100644
--- a/apex/contrib/torchsched/config.py
+++ b/apex/contrib/torchsched/config.py
@@ -71,7 +71,7 @@ def __get_dump_code_backends_and_dir(
     dump_code_dir,
 ) = __get_dump_code_backends_and_dir(os.getenv("TORCH_SCHED_DUMP_CODE"))
 
-from torch.utils._config_module import install_config_module  # noqa: E402
+from torch.utils._config_module import install_config_module
 
 # adds patch, save_config, etc
 install_config_module(sys.modules[__name__])
diff --git a/apex/contrib/torchsched/inductor/event.py b/apex/contrib/torchsched/inductor/event.py
index f6fc9878f..79befd3af 100644
--- a/apex/contrib/torchsched/inductor/event.py
+++ b/apex/contrib/torchsched/inductor/event.py
@@ -15,14 +15,15 @@
 import functools
 import itertools
 
-from torch._inductor.codegen.wrapper import IndentedBuffer
-from torch._inductor.codegen.wrapper import WrapperLine
+from torch._inductor.codegen.wrapper import IndentedBuffer, WrapperLine
 
 import apex.contrib.torchsched.config as torchsched_config
-from apex.contrib.torchsched.inductor._utils import DEFAULT_STREAM_IDX
-from apex.contrib.torchsched.inductor._utils import ENTRANCE_EVENT
-from apex.contrib.torchsched.inductor._utils import EVENT_NAME_TEMPLATE
-from apex.contrib.torchsched.inductor._utils import get_stream_name
+from apex.contrib.torchsched.inductor._utils import (
+    DEFAULT_STREAM_IDX,
+    ENTRANCE_EVENT,
+    EVENT_NAME_TEMPLATE,
+    get_stream_name,
+)
 
 
 @functools.total_ordering
diff --git a/apex/contrib/torchsched/inductor/graph.py b/apex/contrib/torchsched/inductor/graph.py
index 8596f99bd..877ec0192 100644
--- a/apex/contrib/torchsched/inductor/graph.py
+++ b/apex/contrib/torchsched/inductor/graph.py
@@ -7,9 +7,11 @@
 from typing import TYPE_CHECKING
 
 import torch
-from torch._inductor.codegen.common import get_scheduling_for_device
-from torch._inductor.codegen.common import get_wrapper_codegen_for_device
-from torch._inductor.codegen.common import register_backend_for_device
+from torch._inductor.codegen.common import (
+    get_scheduling_for_device,
+    get_wrapper_codegen_for_device,
+    register_backend_for_device,
+)
 from torch._inductor.codegen.wrapper import PythonWrapperCodegen
 from torch._inductor.graph import GraphLowering
 from torch._inductor.scheduler import Scheduler
diff --git a/apex/contrib/torchsched/inductor/scheduler.py b/apex/contrib/torchsched/inductor/scheduler.py
index e0ef61d39..4c6fec492 100644
--- a/apex/contrib/torchsched/inductor/scheduler.py
+++ b/apex/contrib/torchsched/inductor/scheduler.py
@@ -5,28 +5,27 @@
 import collections
 import itertools
 import re
-from typing import TYPE_CHECKING
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import torch
 import torch._inductor.config as inductor_config
 from torch._inductor import ir
 from torch._inductor.dependencies import WeakDep
-from torch._inductor.scheduler import BaseSchedulerNode
-from torch._inductor.scheduler import ExternKernelSchedulerNode
-from torch._inductor.scheduler import ForeachKernelSchedulerNode
-from torch._inductor.scheduler import FusedSchedulerNode
-from torch._inductor.scheduler import NopKernelSchedulerNode
-from torch._inductor.scheduler import Scheduler
-from torch._inductor.scheduler import SchedulerNode
+from torch._inductor.scheduler import (
+    BaseSchedulerNode,
+    ExternKernelSchedulerNode,
+    ForeachKernelSchedulerNode,
+    FusedSchedulerNode,
+    NopKernelSchedulerNode,
+    Scheduler,
+    SchedulerNode,
+)
 from torch._inductor.utils import device_need_guard
 from torch._inductor.virtualized import V
 
 from apex.contrib.torchsched import config
-from apex.contrib.torchsched.inductor._utils import DEFAULT_STREAM_IDX
-from apex.contrib.torchsched.inductor._utils import get_stream_name
-from apex.contrib.torchsched.inductor.event import CudaEventFactory
-from apex.contrib.torchsched.inductor.event import CudaEventSym
+from apex.contrib.torchsched.inductor._utils import DEFAULT_STREAM_IDX, get_stream_name
+from apex.contrib.torchsched.inductor.event import CudaEventFactory, CudaEventSym
 from apex.contrib.torchsched.inductor.wrapper import EnterCudaStreamContextLine
 
 if TYPE_CHECKING:
diff --git a/apex/contrib/torchsched/inductor/wrapper.py b/apex/contrib/torchsched/inductor/wrapper.py
index ef6a2f3eb..40cf189fa 100644
--- a/apex/contrib/torchsched/inductor/wrapper.py
+++ b/apex/contrib/torchsched/inductor/wrapper.py
@@ -14,19 +14,23 @@
 import dataclasses
 from typing import TYPE_CHECKING
 
-from torch._inductor.codegen.wrapper import EnterDeviceContextManagerLine
-from torch._inductor.codegen.wrapper import ExitDeviceContextManagerLine
-from torch._inductor.codegen.wrapper import IndentedBuffer
-from torch._inductor.codegen.wrapper import PythonWrapperCodegen
-from torch._inductor.codegen.wrapper import SubgraphPythonWrapperCodegen
-from torch._inductor.codegen.wrapper import WrapperLine
+from torch._inductor.codegen.wrapper import (
+    EnterDeviceContextManagerLine,
+    ExitDeviceContextManagerLine,
+    IndentedBuffer,
+    PythonWrapperCodegen,
+    SubgraphPythonWrapperCodegen,
+    WrapperLine,
+)
 from torch._inductor.virtualized import V
 
-import apex.contrib.torchsched.config as config
-from apex.contrib.torchsched.inductor._utils import DEFAULT_STREAM
-from apex.contrib.torchsched.inductor._utils import ENTRANCE_EVENT
-from apex.contrib.torchsched.inductor._utils import STREAM_NAME_TEMPLATE
-from apex.contrib.torchsched.inductor._utils import get_stream_name
+from apex.contrib.torchsched import config
+from apex.contrib.torchsched.inductor._utils import (
+    DEFAULT_STREAM,
+    ENTRANCE_EVENT,
+    STREAM_NAME_TEMPLATE,
+    get_stream_name,
+)
 
 if TYPE_CHECKING:
     from torch._inductor.graph import GraphLowering
diff --git a/apex/contrib/torchsched/ops/layer_norm.py b/apex/contrib/torchsched/ops/layer_norm.py
index c58a4b13a..6acc7e4bb 100644
--- a/apex/contrib/torchsched/ops/layer_norm.py
+++ b/apex/contrib/torchsched/ops/layer_norm.py
@@ -277,7 +277,7 @@ def layer_norm(
     #   * Shape (N, S, H), normalized_shape (H,);
     #   * Shape (N, C, H, W), normalized_shape (C, H, W);
     # cuDNN LayerNorm expects shape (M, N, 1, 1) and normalized_shape (1, N, 1, 1)
-    if tuple(x.shape[-len(normalized_shape) :]) != tuple(normalized_shape):  # noqa: E203
+    if tuple(x.shape[-len(normalized_shape) :]) != tuple(normalized_shape):
         raise ValueError(
             f"CuDNN LayerNorm expects `x.shape[{-len(normalized_shape)}:]` equals to "
             f"`normalized_shape`, but got:\n    {x.shape=}, {normalized_shape=}",
diff --git a/apex/contrib/torchsched/passes/pre_grad_passes.py b/apex/contrib/torchsched/passes/pre_grad_passes.py
index 6b12526db..fc112dcd5 100644
--- a/apex/contrib/torchsched/passes/pre_grad_passes.py
+++ b/apex/contrib/torchsched/passes/pre_grad_passes.py
@@ -10,8 +10,7 @@
 from torch.fx import replace_pattern
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
-    from collections.abc import Sequence
+    from collections.abc import Callable, Sequence
 
 from apex.contrib.torchsched import config
 
diff --git a/apex/contrib/transducer/__init__.py b/apex/contrib/transducer/__init__.py
index 955ca1808..f6adb353c 100755
--- a/apex/contrib/transducer/__init__.py
+++ b/apex/contrib/transducer/__init__.py
@@ -1,3 +1,2 @@
-from .transducer import TransducerJoint
-from .transducer import TransducerLoss
 from . import _transducer_ref
+from .transducer import TransducerJoint, TransducerLoss
diff --git a/apex/contrib/transducer/transducer.py b/apex/contrib/transducer/transducer.py
index bb53d39a6..41b660f0f 100755
--- a/apex/contrib/transducer/transducer.py
+++ b/apex/contrib/transducer/transducer.py
@@ -1,6 +1,6 @@
 import torch
-import transducer_loss_cuda
 import transducer_joint_cuda
+import transducer_loss_cuda
 
 
 class TransducerJoint(torch.nn.Module):
@@ -35,7 +35,7 @@ def __init__(
         dropout_prob=0,
         probe_mask=False,
     ):
-        super(TransducerJoint, self).__init__()
+        super().__init__()
         self.pack_output = pack_output
         self.relu = relu
         self.dropout = dropout
@@ -100,7 +100,7 @@ class TransducerLoss(torch.nn.Module):
     """
 
     def __init__(self, fuse_softmax_backward=True, opt=1, packed_input=False):
-        super(TransducerLoss, self).__init__()
+        super().__init__()
         self.fuse_softmax_backward = fuse_softmax_backward
         self.opt = opt
         self.packed_input = packed_input
diff --git a/apex/contrib/xentropy/__init__.py b/apex/contrib/xentropy/__init__.py
index 4c8cbeeea..4367cae49 100644
--- a/apex/contrib/xentropy/__init__.py
+++ b/apex/contrib/xentropy/__init__.py
@@ -1,6 +1,5 @@
 from .softmax_xentropy import SoftmaxCrossEntropyLoss
 
-
 __all__ = [
     "SoftmaxCrossEntropyLoss",
 ]
diff --git a/apex/contrib/xentropy/softmax_xentropy.py b/apex/contrib/xentropy/softmax_xentropy.py
index 528f743b0..d5d94c68f 100644
--- a/apex/contrib/xentropy/softmax_xentropy.py
+++ b/apex/contrib/xentropy/softmax_xentropy.py
@@ -1,5 +1,4 @@
 import torch
-
 import xentropy_cuda
 
 
diff --git a/apex/distributed_testing/distributed_test_base.py b/apex/distributed_testing/distributed_test_base.py
index 8791777fb..b54ca6ef2 100644
--- a/apex/distributed_testing/distributed_test_base.py
+++ b/apex/distributed_testing/distributed_test_base.py
@@ -1,13 +1,12 @@
 import os
 import sys
 import unittest
-from packaging.version import Version, parse
 
 import torch
+from packaging.version import Version, parse
 from torch import distributed as dist
+from torch.testing._internal import common_distributed, common_utils
 from torch.utils import collect_env
-from torch.testing._internal import common_utils
-from torch.testing._internal import common_distributed
 
 from apex.distributed_testing._ucc_util import HAS_UCC
 
diff --git a/apex/fused_dense/fused_dense.py b/apex/fused_dense/fused_dense.py
index 239d12727..ff8c073b9 100644
--- a/apex/fused_dense/fused_dense.py
+++ b/apex/fused_dense/fused_dense.py
@@ -1,6 +1,7 @@
+import fused_dense_cuda
 import torch
 from torch import nn
-import fused_dense_cuda
+
 from apex._autocast_utils import _cast_if_autocast_enabled
 
 
@@ -77,7 +78,7 @@ def _fused_dense_gelu_dense(input, weight1, bias1, weight2, bias2):
 
 class FusedDense(nn.Module):
     def __init__(self, in_features, out_features, bias=True):
-        super(FusedDense, self).__init__()
+        super().__init__()
         self.in_features = in_features
         self.out_features = out_features
         self.weight = nn.Parameter(torch.empty(out_features, in_features))
@@ -96,7 +97,7 @@ def forward(self, input):
 
 class FusedDenseGeluDense(nn.Module):
     def __init__(self, in_features, intermediate_features, out_features, bias=True):
-        super(FusedDenseGeluDense, self).__init__()
+        super().__init__()
         assert bias == True, "DenseGeluDense module without bias is currently not supported"
         self.in_features = in_features
         self.intermediate_features = intermediate_features
diff --git a/apex/mlp/mlp.py b/apex/mlp/mlp.py
index 4297b0a73..ab81f2282 100644
--- a/apex/mlp/mlp.py
+++ b/apex/mlp/mlp.py
@@ -1,11 +1,11 @@
-from copy import copy
 import math
+from copy import copy
 
+import mlp_cuda
 import torch
 from torch import nn
 
 from apex._autocast_utils import _cast_if_autocast_enabled
-import mlp_cuda
 
 
 class MlpFunction(torch.autograd.Function):
@@ -59,12 +59,12 @@ def __init__(self, mlp_sizes, bias=True, activation="relu"):
         for i in range(self.num_layers):
             w = torch.nn.Parameter(torch.empty(mlp_sizes[i + 1], mlp_sizes[i]))
             self.weights.append(w)
-            name = "weight_{}".format(i)
+            name = f"weight_{i}"
             setattr(self, name, w)
             if self.bias:
                 b = torch.nn.Parameter(torch.empty(mlp_sizes[i + 1]))
                 self.biases.append(b)
-                name = "bias_{}".format(i)
+                name = f"bias_{i}"
                 setattr(self, name, b)
 
         self.reset_parameters()
diff --git a/apex/multi_tensor_apply/multi_tensor_apply.py b/apex/multi_tensor_apply/multi_tensor_apply.py
index ba2c21ada..63c212f7d 100644
--- a/apex/multi_tensor_apply/multi_tensor_apply.py
+++ b/apex/multi_tensor_apply/multi_tensor_apply.py
@@ -1,4 +1,4 @@
-class MultiTensorApply(object):
+class MultiTensorApply:
     available = False
     warned = False
 
diff --git a/apex/normalization/__init__.py b/apex/normalization/__init__.py
index 510a32be8..a47775d51 100644
--- a/apex/normalization/__init__.py
+++ b/apex/normalization/__init__.py
@@ -1,6 +1,6 @@
 from .fused_layer_norm import (
     FusedLayerNorm,
-    MixedFusedLayerNorm,
     FusedRMSNorm,
+    MixedFusedLayerNorm,
     MixedFusedRMSNorm,
 )
diff --git a/apex/normalization/fused_layer_norm.py b/apex/normalization/fused_layer_norm.py
index a0f3833bc..15f329e91 100644
--- a/apex/normalization/fused_layer_norm.py
+++ b/apex/normalization/fused_layer_norm.py
@@ -1,11 +1,11 @@
 import importlib
 import numbers
+from typing import List, Tuple
 
 import torch
-from torch.nn.parameter import Parameter
-from torch.nn import init
 from torch.nn import functional as F
-from typing import List, Tuple
+from torch.nn import init
+from torch.nn.parameter import Parameter
 
 from apex._autocast_utils import _cast_if_autocast_enabled
 
@@ -81,10 +81,10 @@ def fused_layer_norm_affine_fwd(
         input: torch.Tensor,
         weight: torch.Tensor,
         bias: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         global fused_layer_norm_cuda
         if fused_layer_norm_cuda is None:
             fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
@@ -102,10 +102,10 @@ def fused_layer_norm_affine_fwd_fake(
         input: torch.Tensor,
         weight: torch.Tensor,
         bias: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         input = input.contiguous()
         weight = weight.contiguous()
         bias = bias.contiguous()
@@ -127,12 +127,12 @@ def fused_layer_norm_affine_bwd(
         mean: torch.Tensor,
         invvar: torch.Tensor,
         input_or_output: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         weight: torch.Tensor,
         bias: torch.Tensor,
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         grad_input, grad_weight, grad_bias = fused_layer_norm_cuda.backward_affine(
             grad_output.contiguous(),
             mean,
@@ -152,12 +152,12 @@ def fused_layer_norm_affine_bwd_fake(
         mean: torch.Tensor,
         invvar: torch.Tensor,
         input_or_output: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         weight: torch.Tensor,
         bias: torch.Tensor,
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         grad_input = torch.empty_like(input_or_output)
         grad_weight = torch.empty_like(weight)
         grad_bias = torch.empty_like(bias)
@@ -241,10 +241,10 @@ def backward(ctx, grad_output):
     def fused_rms_norm_affine_fwd(
         input: torch.Tensor,
         weight: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         global fused_layer_norm_cuda
         if fused_layer_norm_cuda is None:
             fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
@@ -260,10 +260,10 @@ def fused_rms_norm_affine_fwd(
     def fused_rms_norm_affine_fwd_fake(
         input: torch.Tensor,
         weight: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         input = input.contiguous()
         weight = weight.contiguous()
         idiff = input.ndim - len(normalized_shape)
@@ -290,11 +290,11 @@ def fused_rms_norm_affine_bwd(
         grad_output: torch.Tensor,
         invvar: torch.Tensor,
         input_or_output: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         weight: torch.Tensor,
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         grad_input, grad_weight = fused_layer_norm_cuda.rms_backward_affine(
             grad_output.contiguous(),
             invvar,
@@ -311,11 +311,11 @@ def fused_rms_norm_affine_bwd_fake(
         grad_output: torch.Tensor,
         invvar: torch.Tensor,
         input_or_output: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         weight: torch.Tensor,
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         grad_input = torch.empty_like(input_or_output)
         grad_weight = torch.empty_like(weight)
         return grad_input, grad_weight
@@ -433,10 +433,10 @@ def backward(ctx, grad_output):
     @torch.library.custom_op("apex::fused_layer_norm_fwd", mutates_args=())
     def fused_layer_norm_fwd(
         input: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         global fused_layer_norm_cuda
         if fused_layer_norm_cuda is None:
             fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
@@ -448,10 +448,10 @@ def fused_layer_norm_fwd(
     @fused_layer_norm_fwd.register_fake
     def fused_layer_norm_fwd_fake(
         input: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         input = input.contiguous()
         idiff = input.ndim - len(normalized_shape)
         n = 1
@@ -471,7 +471,7 @@ def fused_layer_norm_bwd(
         mean: torch.Tensor,
         invvar: torch.Tensor,
         input_or_output: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
     ) -> torch.Tensor:
@@ -492,7 +492,7 @@ def fused_layer_norm_bwd_fake(
         mean: torch.Tensor,
         invvar: torch.Tensor,
         input_or_output: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
     ) -> torch.Tensor:
@@ -567,10 +567,10 @@ def backward(ctx, grad_output):
     @torch.library.custom_op("apex::fused_rms_norm_fwd", mutates_args=())
     def fused_rms_norm_fwd(
         input: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         global fused_layer_norm_cuda
         if fused_layer_norm_cuda is None:
             fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
@@ -582,10 +582,10 @@ def fused_rms_norm_fwd(
     @fused_rms_norm_fwd.register_fake
     def fused_rms_norm_fwd_fake(
         input: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         input = input.contiguous()
         idiff = input.ndim - len(normalized_shape)
         n = 1
@@ -611,7 +611,7 @@ def fused_rms_norm_bwd(
         grad_output: torch.Tensor,
         invvar: torch.Tensor,
         input_or_output: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
     ) -> torch.Tensor:
@@ -630,7 +630,7 @@ def fused_rms_norm_bwd_fake(
         grad_output: torch.Tensor,
         invvar: torch.Tensor,
         input_or_output: torch.Tensor,
-        normalized_shape: List[int],
+        normalized_shape: list[int],
         eps: float,
         memory_efficient: bool = False,
     ) -> torch.Tensor:
diff --git a/apex/optimizers/__init__.py b/apex/optimizers/__init__.py
index 25c178c5f..a4127a654 100644
--- a/apex/optimizers/__init__.py
+++ b/apex/optimizers/__init__.py
@@ -1,6 +1,6 @@
-from .fused_sgd import FusedSGD
+from .fused_adagrad import FusedAdagrad
 from .fused_adam import FusedAdam
-from .fused_novograd import FusedNovoGrad
 from .fused_lamb import FusedLAMB
-from .fused_adagrad import FusedAdagrad
 from .fused_mixed_precision_lamb import FusedMixedPrecisionLamb
+from .fused_novograd import FusedNovoGrad
+from .fused_sgd import FusedSGD
diff --git a/apex/optimizers/fused_adagrad.py b/apex/optimizers/fused_adagrad.py
index d9dbcd7cc..b5133d250 100644
--- a/apex/optimizers/fused_adagrad.py
+++ b/apex/optimizers/fused_adagrad.py
@@ -1,4 +1,5 @@
 import torch
+
 from apex.multi_tensor_apply import multi_tensor_applier
 
 
@@ -51,7 +52,7 @@ def __init__(
         adagrad_w_mode=False,
     ):
         defaults = dict(lr=lr, eps=eps, weight_decay=weight_decay)
-        super(FusedAdagrad, self).__init__(params, defaults)
+        super().__init__(params, defaults)
         self.adagrad_w_mode = 1 if adagrad_w_mode else 0
         self.set_grad_none = set_grad_none
 
@@ -70,7 +71,7 @@ def zero_grad(self):
                 for p in group["params"]:
                     p.grad = None
         else:
-            super(FusedAdagrad, self).zero_grad()
+            super().zero_grad()
 
     def step(self, closure=None):
         """Performs a single optimization step.
diff --git a/apex/optimizers/fused_adam.py b/apex/optimizers/fused_adam.py
index 45d7e017e..7f063df0d 100644
--- a/apex/optimizers/fused_adam.py
+++ b/apex/optimizers/fused_adam.py
@@ -1,4 +1,5 @@
 import torch
+
 from apex.multi_tensor_apply import multi_tensor_applier
 
 
@@ -94,7 +95,7 @@ def __init__(
             eps=eps,
             weight_decay=weight_decay,
         )
-        super(FusedAdam, self).__init__(params, defaults)
+        super().__init__(params, defaults)
         self.adam_w_mode = 1 if adam_w_mode else 0
         self.set_grad_none = set_grad_none
 
@@ -141,7 +142,7 @@ def zero_grad(self):
                 for p in group["params"]:
                     p.grad = None
         else:
-            super(FusedAdam, self).zero_grad()
+            super().zero_grad()
 
     def step(
         self,
diff --git a/apex/optimizers/fused_lamb.py b/apex/optimizers/fused_lamb.py
index a5630a03c..86ddba6ab 100644
--- a/apex/optimizers/fused_lamb.py
+++ b/apex/optimizers/fused_lamb.py
@@ -1,4 +1,5 @@
 import torch
+
 from apex.multi_tensor_apply import multi_tensor_applier
 
 
@@ -86,7 +87,7 @@ def __init__(
             grad_averaging=grad_averaging,
             max_grad_norm=max_grad_norm,
         )
-        super(FusedLAMB, self).__init__(params, defaults)
+        super().__init__(params, defaults)
         if multi_tensor_applier.available:
             import amp_C
 
@@ -109,7 +110,7 @@ def zero_grad(self):
                 for p in group["params"]:
                     p.grad = None
         else:
-            super(FusedLAMB, self).zero_grad()
+            super().zero_grad()
 
     def step(self, closure=None):
         """Performs a single optimization step.
diff --git a/apex/optimizers/fused_mixed_precision_lamb.py b/apex/optimizers/fused_mixed_precision_lamb.py
index 2ecdddfd2..f6302a536 100644
--- a/apex/optimizers/fused_mixed_precision_lamb.py
+++ b/apex/optimizers/fused_mixed_precision_lamb.py
@@ -1,7 +1,9 @@
-import torch
+from collections import abc as container_abcs
+from collections import defaultdict
 from copy import deepcopy
 from itertools import chain
-from collections import defaultdict, abc as container_abcs
+
+import torch
 
 from apex.multi_tensor_apply import multi_tensor_applier
 
@@ -39,7 +41,7 @@ def __init__(
         )
 
         # init base module
-        super(FusedMixedPrecisionLamb, self).__init__(params, defaults)
+        super().__init__(params, defaults)
 
         # The learning rate (lr) and optimizer step (step) should be located on device
         # in order to faciliated device sync free execution
@@ -97,8 +99,8 @@ def load_state_dict(self, state_dict):
         id_map = {
             old_id: p
             for old_id, p in zip(
-                chain.from_iterable((g["params"] for g in saved_groups)),
-                chain.from_iterable((g["params"] for g in groups)),
+                chain.from_iterable(g["params"] for g in saved_groups),
+                chain.from_iterable(g["params"] for g in groups),
             )
         }
 
diff --git a/apex/optimizers/fused_novograd.py b/apex/optimizers/fused_novograd.py
index b72e2e3e6..bf938c27a 100644
--- a/apex/optimizers/fused_novograd.py
+++ b/apex/optimizers/fused_novograd.py
@@ -1,4 +1,5 @@
 import torch
+
 from apex.multi_tensor_apply import multi_tensor_applier
 
 
@@ -91,7 +92,7 @@ def __init__(
             norm_type=norm_type,
             init_zero=init_zero,
         )
-        super(FusedNovoGrad, self).__init__(params, defaults)
+        super().__init__(params, defaults)
         if multi_tensor_applier.available:
             import amp_C
             # Skip buffer
@@ -113,10 +114,10 @@ def zero_grad(self):
                 for p in group["params"]:
                     p.grad = None
         else:
-            super(FusedNovoGrad, self).zero_grad()
+            super().zero_grad()
 
     def load_state_dict(self, state_dict):
-        super(FusedNovoGrad, self).load_state_dict(state_dict)
+        super().load_state_dict(state_dict)
         # in case exp_avg_sq is not on the same device as params, move it there
         for group in self.param_groups:
             if len(group["params"]) > 0:
diff --git a/apex/optimizers/fused_sgd.py b/apex/optimizers/fused_sgd.py
index f4cec9f57..4d7d05ee8 100644
--- a/apex/optimizers/fused_sgd.py
+++ b/apex/optimizers/fused_sgd.py
@@ -87,11 +87,11 @@ def __init__(
         set_grad_none=False,
     ):
         if lr is not required and lr < 0.0:
-            raise ValueError("Invalid learning rate: {}".format(lr))
+            raise ValueError(f"Invalid learning rate: {lr}")
         if momentum < 0.0:
-            raise ValueError("Invalid momentum value: {}".format(momentum))
+            raise ValueError(f"Invalid momentum value: {momentum}")
         if weight_decay < 0.0:
-            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
 
         defaults = dict(
             lr=lr,
@@ -102,7 +102,7 @@ def __init__(
         )
         if nesterov and (momentum <= 0 or dampening != 0):
             raise ValueError("Nesterov momentum requires a momentum and zero dampening")
-        super(FusedSGD, self).__init__(params, defaults)
+        super().__init__(params, defaults)
 
         self.wd_after_momentum = wd_after_momentum
         self.materialize_master_grads = materialize_master_grads
@@ -122,7 +122,7 @@ def __init__(
             raise RuntimeError("apex.optimizers.FusedSGD requires cuda extensions")
 
     def __setstate__(self, state):
-        super(FusedSGD, self).__setstate__(state)
+        super().__setstate__(state)
         for group in self.param_groups:
             group.setdefault("nesterov", False)
 
@@ -132,7 +132,7 @@ def zero_grad(self):
                 for p in group["params"]:
                     p.grad = None
         else:
-            super(FusedSGD, self).zero_grad()
+            super().zero_grad()
 
     def get_momentums(self, params):
         momentums = []
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 208a20ed0..ccf3e0050 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -1,5 +1,4 @@
 #!/usr/bin/env python3
-# -*- coding: utf-8 -*-
 #
 # PyTorch documentation build configuration file, created by
 # sphinx-quickstart on Fri Dec 23 13:31:47 2016.
@@ -25,7 +24,6 @@
 # import multiproc
 import sphinx_rtd_theme
 
-
 # -- General configuration ------------------------------------------------
 
 # If your documentation needs a minimal Sphinx version, state it here.
@@ -198,8 +196,8 @@
 # See http://stackoverflow.com/a/41184353/3343043
 
 from docutils import nodes
-from sphinx.util.docfields import TypedField
 from sphinx import addnodes
+from sphinx.util.docfields import TypedField
 
 
 def patched_make_field(self, types, domain, items, **kw):
diff --git a/examples/dcgan/main_amp.py b/examples/dcgan/main_amp.py
index be1a2894f..7476a6ff9 100644
--- a/examples/dcgan/main_amp.py
+++ b/examples/dcgan/main_amp.py
@@ -1,16 +1,15 @@
-from __future__ import print_function
 import argparse
 import os
 import random
+
 import torch
-import torch.nn as nn
 import torch.nn.parallel
-import torch.backends.cudnn as cudnn
-import torch.optim as optim
 import torch.utils.data
 import torchvision.datasets as dset
-import torchvision.transforms as transforms
 import torchvision.utils as vutils
+from torch import nn, optim
+from torch.backends import cudnn
+from torchvision import transforms
 
 try:    
     from apex import amp
@@ -122,7 +121,7 @@ def weights_init(m):
 
 class Generator(nn.Module):
     def __init__(self, ngpu):
-        super(Generator, self).__init__()
+        super().__init__()
         self.ngpu = ngpu
         self.main = nn.Sequential(
             # input is Z, going into a convolution
@@ -164,7 +163,7 @@ def forward(self, input):
 
 class Discriminator(nn.Module):
     def __init__(self, ngpu):
-        super(Discriminator, self).__init__()
+        super().__init__()
         self.ngpu = ngpu
         self.main = nn.Sequential(
             # input is (nc) x 64 x 64
diff --git a/examples/imagenet/main_amp.py b/examples/imagenet/main_amp.py
index 384e5bda9..236d69732 100644
--- a/examples/imagenet/main_amp.py
+++ b/examples/imagenet/main_amp.py
@@ -3,21 +3,18 @@
 import shutil
 import time
 
+import numpy as np
 import torch
-import torch.nn as nn
-import torch.nn.parallel
-import torch.backends.cudnn as cudnn
 import torch.distributed as dist
+import torch.nn.parallel
 import torch.optim
 import torch.utils.data
 import torch.utils.data.distributed
-import torchvision.transforms as transforms
-import torchvision.datasets as datasets
-import torchvision.models as models
-
-import numpy as np
-
+from torch import nn
+from torch.backends import cudnn
 from torch.nn.parallel import DistributedDataParallel as DDP
+from torchvision import datasets, models, transforms
+
 
 def to_python_float(scalar_tensor: torch.Tensor):
     return scalar_tensor.float().item()
@@ -93,11 +90,11 @@ def main():
     global best_prec1, args
 
     args = parse()
-    print("opt_level = {}".format(args.opt_level))
-    print("keep_batchnorm_fp32 = {}".format(args.keep_batchnorm_fp32), type(args.keep_batchnorm_fp32))
-    print("loss_scale = {}".format(args.loss_scale), type(args.loss_scale))
+    print(f"opt_level = {args.opt_level}")
+    print(f"keep_batchnorm_fp32 = {args.keep_batchnorm_fp32}", type(args.keep_batchnorm_fp32))
+    print(f"loss_scale = {args.loss_scale}", type(args.loss_scale))
 
-    print("\nCUDNN VERSION: {}\n".format(torch.backends.cudnn.version()))
+    print(f"\nCUDNN VERSION: {torch.backends.cudnn.version()}\n")
 
     cudnn.benchmark = True
     best_prec1 = 0
@@ -130,10 +127,10 @@ def main():
 
     # create model
     if args.pretrained:
-        print("=> using pre-trained model '{}'".format(args.arch))
+        print(f"=> using pre-trained model '{args.arch}'")
         model = models.__dict__[args.arch](pretrained=True)
     else:
-        print("=> creating model '{}'".format(args.arch))
+        print(f"=> creating model '{args.arch}'")
         model = models.__dict__[args.arch]()
 
     if args.sync_bn:
@@ -161,7 +158,7 @@ def main():
         # Use a local scope to avoid dangling references
         def resume():
             if os.path.isfile(args.resume):
-                print("=> loading checkpoint '{}'".format(args.resume))
+                print(f"=> loading checkpoint '{args.resume}'")
                 checkpoint = torch.load(args.resume, map_location = lambda storage, loc: storage.cuda(args.gpu))
                 args.start_epoch = checkpoint['epoch']
                 global best_prec1
@@ -171,7 +168,7 @@ def resume():
                 print("=> loaded checkpoint '{}' (epoch {})"
                       .format(args.resume, checkpoint['epoch']))
             else:
-                print("=> no checkpoint found at '{}'".format(args.resume))
+                print(f"=> no checkpoint found at '{args.resume}'")
         resume()
 
     # Data loading code
@@ -244,7 +241,7 @@ def resume():
                 'optimizer' : optimizer.state_dict(),
             }, is_best)
 
-class data_prefetcher():
+class data_prefetcher:
     def __init__(self, loader):
         self.loader = iter(loader)
         self.stream = torch.cuda.Stream()
@@ -315,10 +312,10 @@ def train(train_loader, model, criterion, optimizer, scaler, epoch):
     while input is not None:
         i += 1
         if args.prof >= 0 and i == args.prof:
-            print("Profiling begun at iteration {}".format(i))
+            print(f"Profiling begun at iteration {i}")
             torch.cuda.cudart().cudaProfilerStart()
 
-        if args.prof >= 0: torch.cuda.nvtx.range_push("Body of iteration {}".format(i))
+        if args.prof >= 0: torch.cuda.nvtx.range_push(f"Body of iteration {i}")
 
         adjust_learning_rate(optimizer, epoch, i, len(train_loader))
 
@@ -370,17 +367,12 @@ def train(train_loader, model, criterion, optimizer, scaler, epoch):
             end = time.time()
 
             if args.local_rank == 0:
-                print('Epoch: [{0}][{1}/{2}]\t'
-                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
-                      'Speed {3:.3f} ({4:.3f})\t'
-                      'Loss {loss.val:.10f} ({loss.avg:.4f})\t'
-                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
-                      'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
-                       epoch, i, len(train_loader),
-                       args.world_size*args.batch_size/batch_time.val,
-                       args.world_size*args.batch_size/batch_time.avg,
-                       batch_time=batch_time,
-                       loss=losses, top1=top1, top5=top5))
+                print(f'Epoch: [{epoch}][{i}/{len(train_loader)}]\t'
+                      f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                      f'Speed {args.world_size*args.batch_size/batch_time.val:.3f} ({args.world_size*args.batch_size/batch_time.avg:.3f})\t'
+                      f'Loss {losses.val:.10f} ({losses.avg:.4f})\t'
+                      f'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
+                      f'Prec@5 {top5.val:.3f} ({top5.avg:.3f})')
         if args.prof >= 0: torch.cuda.nvtx.range_push("prefetcher.next()")
         input, target = prefetcher.next()
         if args.prof >= 0: torch.cuda.nvtx.range_pop()
@@ -389,7 +381,7 @@ def train(train_loader, model, criterion, optimizer, scaler, epoch):
         if args.prof >= 0: torch.cuda.nvtx.range_pop()
 
         if args.prof >= 0 and i == args.prof + 10:
-            print("Profiling ended at iteration {}".format(i))
+            print(f"Profiling ended at iteration {i}")
             torch.cuda.cudart().cudaProfilerStop()
             quit()
 
@@ -436,22 +428,17 @@ def validate(val_loader, model, criterion):
 
         # TODO:  Change timings to mirror train().
         if args.local_rank == 0 and i % args.print_freq == 0:
-            print('Test: [{0}/{1}]\t'
-                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
-                  'Speed {2:.3f} ({3:.3f})\t'
-                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
-                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
-                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
-                   i, len(val_loader),
-                   args.world_size * args.batch_size / batch_time.val,
-                   args.world_size * args.batch_size / batch_time.avg,
-                   batch_time=batch_time, loss=losses,
-                   top1=top1, top5=top5))
+            print(f'Test: [{i}/{len(val_loader)}]\t'
+                  f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                  f'Speed {args.world_size * args.batch_size / batch_time.val:.3f} ({args.world_size * args.batch_size / batch_time.avg:.3f})\t'
+                  f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
+                  f'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
+                  f'Prec@5 {top5.val:.3f} ({top5.avg:.3f})')
 
         input, target = prefetcher.next()
 
-    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
-          .format(top1=top1, top5=top5))
+    print(f' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
+          )
 
     return top1.avg
 
@@ -462,7 +449,7 @@ def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
         shutil.copyfile(filename, 'model_best.pth.tar')
 
 
-class AverageMeter(object):
+class AverageMeter:
     """Computes and stores the average and current value"""
     def __init__(self):
         self.reset()
diff --git a/examples/simple/distributed/distributed_data_parallel.py b/examples/simple/distributed/distributed_data_parallel.py
index b364405df..ef401ee2f 100644
--- a/examples/simple/distributed/distributed_data_parallel.py
+++ b/examples/simple/distributed/distributed_data_parallel.py
@@ -1,10 +1,13 @@
-import torch
 import argparse
 import os
-from apex import amp
+
+import torch
+
 # FOR DISTRIBUTED: (can also use torch.nn.parallel.DistributedDataParallel instead)
 from apex.parallel import DistributedDataParallel
 
+from apex import amp
+
 parser = argparse.ArgumentParser()
 # FOR DISTRIBUTED:  Parse for the local_rank argument, which will be supplied
 # automatically by torch.distributed.launch.
diff --git a/setup.py b/setup.py
index 696083f6c..2483979c2 100644
--- a/setup.py
+++ b/setup.py
@@ -1,19 +1,18 @@
-import sys
-import warnings
-import os
-import threading
 import glob
-from packaging.version import parse, Version
-
-from setuptools import setup, find_packages
+import os
 import subprocess
+import sys
+import threading
+import warnings
 
 import torch
+from packaging.version import Version, parse
+from setuptools import find_packages, setup
 from torch.utils.cpp_extension import (
+    CUDA_HOME,
     BuildExtension,
     CppExtension,
     CUDAExtension,
-    CUDA_HOME,
     load,
 )
 
@@ -82,7 +81,7 @@ def check_cuda_torch_binary_vs_bare_metal(cuda_dir):
         raise RuntimeError(
             "Cuda extensions are being compiled with a version of Cuda that does "
             "not match the version used to compile Pytorch binaries.  "
-            "Pytorch binaries were compiled with Cuda {}.\n".format(torch.version.cuda)
+            f"Pytorch binaries were compiled with Cuda {torch.version.cuda}.\n"
             + "In some cases, a minor-version mismatch will not cause later errors:  "
             "https://github.com/NVIDIA/apex/pull/323#discussion_r287021798.  "
             "You can try commenting out this check (at your own risk)."
@@ -140,7 +139,7 @@ def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int
         else:
             os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5"
 
-print("\n\ntorch.__version__  = {}\n\n".format(torch.__version__))
+print(f"\n\ntorch.__version__  = {torch.__version__}\n\n")
 TORCH_MAJOR = int(torch.__version__.split(".")[0])
 TORCH_MINOR = int(torch.__version__.split(".")[1])
 
@@ -157,9 +156,7 @@ def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int
 if "--cpp_ext" in sys.argv or "--cuda_ext" in sys.argv:
     if TORCH_MAJOR == 0:
         raise RuntimeError(
-            "--cpp_ext requires Pytorch 1.0 or later, found torch.__version__ = {}".format(
-                torch.__version__
-            )
+            f"--cpp_ext requires Pytorch 1.0 or later, found torch.__version__ = {torch.__version__}"
         )
 
 if has_flag("--cpp_ext", "APEX_CPP_EXT"):
diff --git a/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py b/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py
index 6779078a3..2f4e27df3 100644
--- a/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py
+++ b/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py
@@ -1,15 +1,11 @@
 import importlib.util
+from itertools import product
 
 import torch
-from apex.normalization import FusedLayerNorm
-from apex.normalization import FusedRMSNorm
-from apex.normalization import MixedFusedLayerNorm
-from apex.normalization import MixedFusedRMSNorm
-
 from torch.testing._internal import common_utils
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 
-from itertools import product
+from apex.normalization import FusedLayerNorm, FusedRMSNorm, MixedFusedLayerNorm, MixedFusedRMSNorm
 
 
 def _prep_inputs(batch_size, normalized_shape, dtype):
diff --git a/tests/L0/run_mlp/test_mlp.py b/tests/L0/run_mlp/test_mlp.py
index 0da9f747c..7d46128bd 100644
--- a/tests/L0/run_mlp/test_mlp.py
+++ b/tests/L0/run_mlp/test_mlp.py
@@ -6,12 +6,11 @@
 import torch
 from torch import nn
 from torch.testing._internal import common_utils
-from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_cuda import tf32_off
+from torch.testing._internal.common_device_type import instantiate_device_type_tests
 
 from apex.mlp import MLP
 
-
 batch_size = 1024
 mlp_sizes = [480, 1024, 1024, 512, 256, 1]
 num_iters = 10
diff --git a/tests/L0/run_optimizers/test_adam.py b/tests/L0/run_optimizers/test_adam.py
index c232fcb50..78b32bb9b 100644
--- a/tests/L0/run_optimizers/test_adam.py
+++ b/tests/L0/run_optimizers/test_adam.py
@@ -15,7 +15,7 @@
 
 class Model(torch.nn.Module):
     def __init__(self):
-        super(Model, self).__init__()
+        super().__init__()
         self.conv1 = nn.Conv2d(1, 6, 5)
         self.relu1 = nn.ReLU()
         self.pool1 = nn.MaxPool2d(2)
diff --git a/tests/L0/run_optimizers/test_fused_novograd.py b/tests/L0/run_optimizers/test_fused_novograd.py
index 894fe51f1..c2b7229b3 100755
--- a/tests/L0/run_optimizers/test_fused_novograd.py
+++ b/tests/L0/run_optimizers/test_fused_novograd.py
@@ -1,10 +1,11 @@
-import torch
-from torch.optim import Optimizer
-import apex
 import unittest
+from itertools import product
 
+import torch
 from test_fused_optimizer import TestFusedOptimizer
-from itertools import product
+from torch.optim import Optimizer
+
+import apex
 
 
 class Novograd(Optimizer):
@@ -37,13 +38,13 @@ def __init__(
         amsgrad=False,
     ):
         if not 0.0 <= lr:
-            raise ValueError("Invalid learning rate: {}".format(lr))
+            raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= eps:
-            raise ValueError("Invalid epsilon value: {}".format(eps))
+            raise ValueError(f"Invalid epsilon value: {eps}")
         if not 0.0 <= betas[0] < 1.0:
-            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
         if not 0.0 <= betas[1] < 1.0:
-            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
         defaults = dict(
             lr=lr,
             betas=betas,
@@ -53,10 +54,10 @@ def __init__(
             amsgrad=amsgrad,
         )
 
-        super(Novograd, self).__init__(params, defaults)
+        super().__init__(params, defaults)
 
     def __setstate__(self, state):
-        super(Novograd, self).__setstate__(state)
+        super().__setstate__(state)
         for group in self.param_groups:
             group.setdefault("amsgrad", False)
 
@@ -129,7 +130,7 @@ def step(self, closure=None):
 
 class TestFusedNovoGrad(TestFusedOptimizer):
     def __init__(self, *args, **kwargs):
-        super(TestFusedNovoGrad, self).__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
 
         # The options for NovoGrad and FusedNovoGrad are very specific if they
         # are expected to behave the same.
diff --git a/tests/L0/run_optimizers/test_fused_optimizer.py b/tests/L0/run_optimizers/test_fused_optimizer.py
index c0993e5bd..806536e6f 100644
--- a/tests/L0/run_optimizers/test_fused_optimizer.py
+++ b/tests/L0/run_optimizers/test_fused_optimizer.py
@@ -1,6 +1,6 @@
-from itertools import product
 import random
 import unittest
+from itertools import product
 
 import torch
 
@@ -53,10 +53,8 @@ def get_max_diff(self, ref_param, tst_param):
             max_abs_diff_p = (p_ref - p_tst).abs().max().item()
             max_rel_diff_p = ((p_ref - p_tst) / p_ref).abs().max().item()
 
-            if max_abs_diff_p > max_abs_diff:
-                max_abs_diff = max_abs_diff_p
-            if max_rel_diff_p > max_rel_diff:
-                max_rel_diff = max_rel_diff_p
+            max_abs_diff = max(max_abs_diff, max_abs_diff_p)
+            max_rel_diff = max(max_rel_diff, max_rel_diff_p)
 
         return max_abs_diff, max_rel_diff
 
@@ -226,7 +224,7 @@ def test_frozen_model(self):
 
 class TestFusedAdagrad(TestFusedOptimizer):
     def __init__(self, *args, **kwargs):
-        super(TestFusedAdagrad, self).__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
         self.options = {"lr": 5e-4, "eps": 1e-08, "weight_decay": 1.0e-5}
         self.ref_optim = torch.optim.Adagrad
         self.fused_optim = apex.optimizers.FusedAdagrad
@@ -294,7 +292,7 @@ def test_adagrad_option(self):
 
 class TestFusedSGD(TestFusedOptimizer):
     def __init__(self, *args, **kwargs):
-        super(TestFusedSGD, self).__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
         self.options = {"lr": 0.25, "momentum": 0.125}
         self.ref_optim = torch.optim.SGD
         self.fused_optim = apex.optimizers.FusedSGD
diff --git a/tests/L0/run_optimizers/test_lamb.py b/tests/L0/run_optimizers/test_lamb.py
index 3b208e61b..664eef452 100644
--- a/tests/L0/run_optimizers/test_lamb.py
+++ b/tests/L0/run_optimizers/test_lamb.py
@@ -1,11 +1,12 @@
-import unittest
 import os
+import unittest
+from itertools import product
 
 import torch
 from torch.optim import Optimizer
+
 import apex
 from apex.multi_tensor_apply import multi_tensor_applier
-from itertools import product
 
 
 class RefLAMB(Optimizer):
@@ -29,15 +30,15 @@ class RefLAMB(Optimizer):
 
     def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.01):
         if not 0.0 <= lr:
-            raise ValueError("Invalid learning rate: {}".format(lr))
+            raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= eps:
-            raise ValueError("Invalid epsilon value: {}".format(eps))
+            raise ValueError(f"Invalid epsilon value: {eps}")
         if not 0.0 <= betas[0] < 1.0:
-            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
         if not 0.0 <= betas[1] < 1.0:
-            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
         defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
-        super(RefLAMB, self).__init__(params, defaults)
+        super().__init__(params, defaults)
         if multi_tensor_applier.available:
             import amp_C
 
diff --git a/tests/L0/run_test.py b/tests/L0/run_test.py
index 7011c37ea..b1ff4e7d8 100644
--- a/tests/L0/run_test.py
+++ b/tests/L0/run_test.py
@@ -13,9 +13,8 @@
 
 import argparse
 import os
-import unittest
 import sys
-
+import unittest
 
 TEST_ROOT = os.path.dirname(os.path.abspath(__file__))
 TEST_DIRS = [
@@ -69,8 +68,9 @@ def main(args: argparse.Namespace) -> None:
 
             warnings.warn("The option of `--xml-report` is deprecated", FutureWarning)
 
+        from datetime import date
+
         import xmlrunner
-        from datetime import date  # NOQA
 
         Runner = xmlrunner.XMLTestRunner
         if args.xml_report:
diff --git a/tests/L1/common/compare.py b/tests/L1/common/compare.py
index 0ed5cfba7..d9ac2416a 100644
--- a/tests/L1/common/compare.py
+++ b/tests/L1/common/compare.py
@@ -1,4 +1,5 @@
 import argparse
+
 import torch
 
 parser = argparse.ArgumentParser(description="Compare")
@@ -39,13 +40,11 @@
 # ugly duplication here...
 if not args.use_baseline:
     for n, (i_e, i_p) in enumerate(zip(dict_e["Iteration"], dict_p["Iteration"])):
-        assert i_e == i_p, "i_e = {}, i_p = {}".format(i_e, i_p)
+        assert i_e == i_p, f"i_e = {i_e}, i_p = {i_p}"
 
         loss_e = dict_e["Loss"][n]
         loss_p = dict_p["Loss"][n]
-        assert loss_e == loss_p, "Iteration {}, loss_e = {}, loss_p = {}".format(
-            i_e, loss_e, loss_p
-        )
+        assert loss_e == loss_p, f"Iteration {i_e}, loss_e = {loss_e}, loss_p = {loss_p}"
         print(
             "{:4} {:15.10f} {:15.10f} {:15.10f} {:15.10f}".format(
                 i_e, loss_e, loss_p, dict_e["Speed"][n], dict_p["Speed"][n]
@@ -53,17 +52,13 @@
         )
 else:
     for n, (i_e, i_p) in enumerate(zip(dict_e["Iteration"], dict_p["Iteration"])):
-        assert i_e == i_p, "i_e = {}, i_p = {}".format(i_e, i_p)
+        assert i_e == i_p, f"i_e = {i_e}, i_p = {i_p}"
 
         loss_e = dict_e["Loss"][n]
         loss_p = dict_p["Loss"][n]
         loss_b = dict_b["Loss"][n]
-        assert loss_e == loss_p, "Iteration {}, loss_e = {}, loss_p = {}".format(
-            i_e, loss_e, loss_p
-        )
-        assert loss_e == loss_b, "Iteration {}, loss_e = {}, loss_b = {}".format(
-            i_e, loss_e, loss_b
-        )
+        assert loss_e == loss_p, f"Iteration {i_e}, loss_e = {loss_e}, loss_p = {loss_p}"
+        assert loss_e == loss_b, f"Iteration {i_e}, loss_e = {loss_e}, loss_b = {loss_b}"
         print(
             "{:4} {:15.10f} {:15.10f} {:15.10f} {:15.10f} {:15.10f} {:15.10f}".format(
                 i_e,
diff --git a/tests/L1/common/main_amp.py b/tests/L1/common/main_amp.py
index 344d9200e..52478df13 100644
--- a/tests/L1/common/main_amp.py
+++ b/tests/L1/common/main_amp.py
@@ -3,23 +3,21 @@
 import shutil
 import time
 
+import numpy as np
 import torch
-import torch.nn as nn
-import torch.nn.parallel
-import torch.backends.cudnn as cudnn
 import torch.distributed as dist
+import torch.nn.parallel
 import torch.optim
 import torch.utils.data
 import torch.utils.data.distributed
-import torchvision.transforms as transforms
-import torchvision.datasets as datasets
-import torchvision.models as models
-
-import numpy as np
+from torch import nn
+from torch.backends import cudnn
+from torchvision import datasets, models, transforms
 
 try:
-    from apex.parallel import DistributedDataParallel as DDP
     from apex.fp16_utils import *
+    from apex.parallel import DistributedDataParallel as DDP
+
     from apex import amp, optimizers
     from apex.multi_tensor_apply import multi_tensor_applier
 except ImportError:
@@ -158,15 +156,15 @@ def fast_collate(batch):
 # that verifies if the backend is what we think it is
 assert multi_tensor_applier.available == args.has_ext
 
-print("opt_level = {}".format(args.opt_level))
+print(f"opt_level = {args.opt_level}")
 print(
-    "keep_batchnorm_fp32 = {}".format(args.keep_batchnorm_fp32),
+    f"keep_batchnorm_fp32 = {args.keep_batchnorm_fp32}",
     type(args.keep_batchnorm_fp32),
 )
-print("loss_scale = {}".format(args.loss_scale), type(args.loss_scale))
+print(f"loss_scale = {args.loss_scale}", type(args.loss_scale))
 
 
-print("\nCUDNN VERSION: {}\n".format(torch.backends.cudnn.version()))
+print(f"\nCUDNN VERSION: {torch.backends.cudnn.version()}\n")
 
 if args.deterministic:
     cudnn.benchmark = False
@@ -195,10 +193,10 @@ def main():
 
     # create model
     if args.pretrained:
-        print("=> using pre-trained model '{}'".format(args.arch))
+        print(f"=> using pre-trained model '{args.arch}'")
         model = models.__dict__[args.arch](pretrained=True)
     else:
-        print("=> creating model '{}'".format(args.arch))
+        print(f"=> creating model '{args.arch}'")
         model = models.__dict__[args.arch]()
 
     if args.sync_bn:
@@ -245,7 +243,7 @@ def main():
         # Use a local scope to avoid dangling references
         def resume():
             if os.path.isfile(args.resume):
-                print("=> loading checkpoint '{}'".format(args.resume))
+                print(f"=> loading checkpoint '{args.resume}'")
                 checkpoint = torch.load(
                     args.resume,
                     map_location=lambda storage, loc: storage.cuda(args.gpu),
@@ -258,7 +256,7 @@ def resume():
                     "=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint["epoch"])
                 )
             else:
-                print("=> no checkpoint found at '{}'".format(args.resume))
+                print(f"=> no checkpoint found at '{args.resume}'")
 
         resume()
 
@@ -461,24 +459,13 @@ def train(train_loader, model, criterion, optimizer, epoch):
         if i % args.print_freq == 0 and i > 1:
             if args.local_rank == 0:
                 print(
-                    "Epoch: [{0}][{1}/{2}]\t"
-                    "Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
-                    "Speed {3:.3f} ({4:.3f})\t"
-                    "Data {data_time.val:.3f} ({data_time.avg:.3f})\t"
-                    "Loss {loss.val:.10f} ({loss.avg:.4f})\t"
-                    "Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t"
-                    "Prec@5 {top5.val:.3f} ({top5.avg:.3f})".format(
-                        epoch,
-                        i,
-                        len(train_loader),
-                        args.world_size * args.batch_size / batch_time.val,
-                        args.world_size * args.batch_size / batch_time.avg,
-                        batch_time=batch_time,
-                        data_time=data_time,
-                        loss=losses,
-                        top1=top1,
-                        top5=top5,
-                    )
+                    f"Epoch: [{epoch}][{i}/{len(train_loader)}]\t"
+                    f"Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
+                    f"Speed {args.world_size * args.batch_size / batch_time.val:.3f} ({args.world_size * args.batch_size / batch_time.avg:.3f})\t"
+                    f"Data {data_time.val:.3f} ({data_time.avg:.3f})\t"
+                    f"Loss {losses.val:.10f} ({losses.avg:.4f})\t"
+                    f"Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t"
+                    f"Prec@5 {top5.val:.3f} ({top5.avg:.3f})"
                 )
             run_info_dict["Iteration"].append(i)
             run_info_dict["Loss"].append(losses.val)
@@ -542,26 +529,17 @@ def validate(val_loader, model, criterion):
 
         if args.local_rank == 0 and i % args.print_freq == 0:
             print(
-                "Test: [{0}/{1}]\t"
-                "Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
-                "Speed {2:.3f} ({3:.3f})\t"
-                "Loss {loss.val:.4f} ({loss.avg:.4f})\t"
-                "Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t"
-                "Prec@5 {top5.val:.3f} ({top5.avg:.3f})".format(
-                    i,
-                    len(val_loader),
-                    args.world_size * args.batch_size / batch_time.val,
-                    args.world_size * args.batch_size / batch_time.avg,
-                    batch_time=batch_time,
-                    loss=losses,
-                    top1=top1,
-                    top5=top5,
-                )
+                f"Test: [{i}/{len(val_loader)}]\t"
+                f"Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
+                f"Speed {args.world_size * args.batch_size / batch_time.val:.3f} ({args.world_size * args.batch_size / batch_time.avg:.3f})\t"
+                f"Loss {losses.val:.4f} ({losses.avg:.4f})\t"
+                f"Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t"
+                f"Prec@5 {top5.val:.3f} ({top5.avg:.3f})"
             )
 
         input, target = prefetcher.next()
 
-    print(" * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}".format(top1=top1, top5=top5))
+    print(f" * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}")
 
     return top1.avg
 
@@ -572,7 +550,7 @@ def save_checkpoint(state, is_best, filename="checkpoint.pth.tar"):
         shutil.copyfile(filename, "model_best.pth.tar")
 
 
-class AverageMeter(object):
+class AverageMeter:
     """Computes and stores the average and current value"""
 
     def __init__(self):
diff --git a/tests/distributed/DDP/ddp_race_condition_test.py b/tests/distributed/DDP/ddp_race_condition_test.py
index 20dd6e24e..c432a5121 100644
--- a/tests/distributed/DDP/ddp_race_condition_test.py
+++ b/tests/distributed/DDP/ddp_race_condition_test.py
@@ -1,10 +1,9 @@
-import torch
-from torch.nn import Parameter
-from torch.nn import Module
-from apex.parallel import DistributedDataParallel as DDP
 import argparse
 import os
 
+import torch
+from apex.parallel import DistributedDataParallel as DDP
+from torch.nn import Module, Parameter
 
 parser = argparse.ArgumentParser(description="allreduce hook example")
 parser.add_argument("--local_rank", default=0, type=int)
@@ -26,7 +25,7 @@
 
 class Model(Module):
     def __init__(self):
-        super(Model, self).__init__()
+        super().__init__()
         self.a = Parameter(torch.cuda.FloatTensor(4096 * 4096).fill_(1.0))
         self.b = Parameter(torch.cuda.FloatTensor(4096 * 4096).fill_(2.0))
 
@@ -55,16 +54,14 @@ def forward(self, input):
 
     # torch.cuda.nvtx.range_push("synchronize() + info")
     # torch.cuda.synchronize()
-    print("i = {}".format(i))
+    print(f"i = {i}")
 
     def info(name, param, val):
         expected = val * 4096 * 4096 * (2.0 * i + 1) / 2.0
         actual = param.grad.data.sum().item()
         print(
             name
-            + ": grad.data_ptr() = {}, expected sum {}, got {}".format(
-                param.grad.data_ptr(), expected, actual
-            )
+            + f": grad.data_ptr() = {param.grad.data_ptr()}, expected sum {expected}, got {actual}"
         )
         return expected == actual
 
diff --git a/tests/distributed/amp_master_params/amp_master_params.py b/tests/distributed/amp_master_params/amp_master_params.py
index 868ff2f41..61c04f817 100644
--- a/tests/distributed/amp_master_params/amp_master_params.py
+++ b/tests/distributed/amp_master_params/amp_master_params.py
@@ -1,11 +1,13 @@
-import torch
 import argparse
 import os
-from apex import amp
+
+import torch
 
 # FOR DISTRIBUTED: (can also use torch.nn.parallel.DistributedDataParallel instead)
 from apex.parallel import DistributedDataParallel
 
+from apex import amp
+
 parser = argparse.ArgumentParser()
 # FOR DISTRIBUTED:  Parse for the local_rank argument, which will be supplied
 # automatically by torch.distributed.launch.
@@ -66,8 +68,8 @@
 if args.local_rank == 0:
     print("final loss = ", loss)
 
-torch.save(list(model.parameters()), "rank{}model.pth".format(torch.distributed.get_rank()))
+torch.save(list(model.parameters()), f"rank{torch.distributed.get_rank()}model.pth")
 torch.save(
     list(amp.master_params(optimizer)),
-    "rank{}master.pth".format(torch.distributed.get_rank()),
+    f"rank{torch.distributed.get_rank()}master.pth",
 )
diff --git a/tests/distributed/synced_batchnorm/python_single_gpu_unit_test.py b/tests/distributed/synced_batchnorm/python_single_gpu_unit_test.py
index e310b9e81..45cacb0c7 100644
--- a/tests/distributed/synced_batchnorm/python_single_gpu_unit_test.py
+++ b/tests/distributed/synced_batchnorm/python_single_gpu_unit_test.py
@@ -1,5 +1,5 @@
-import torch
 import numpy as np
+import torch
 
 
 def compare(desc, inp1, inp2, error):
diff --git a/tests/distributed/synced_batchnorm/single_gpu_unit_test.py b/tests/distributed/synced_batchnorm/single_gpu_unit_test.py
index 18ea55dcc..2bfde0367 100644
--- a/tests/distributed/synced_batchnorm/single_gpu_unit_test.py
+++ b/tests/distributed/synced_batchnorm/single_gpu_unit_test.py
@@ -1,5 +1,6 @@
-import torch
 import numpy as np
+import torch
+
 import apex
 
 if True:
diff --git a/tests/distributed/synced_batchnorm/test_batchnorm1d.py b/tests/distributed/synced_batchnorm/test_batchnorm1d.py
index 360e69441..29a46498f 100644
--- a/tests/distributed/synced_batchnorm/test_batchnorm1d.py
+++ b/tests/distributed/synced_batchnorm/test_batchnorm1d.py
@@ -1,4 +1,5 @@
 import torch
+
 import apex
 
 model = apex.parallel.SyncBatchNorm(4).cuda()
diff --git a/tests/distributed/synced_batchnorm/test_groups.py b/tests/distributed/synced_batchnorm/test_groups.py
index e95aa9984..5bec405fd 100644
--- a/tests/distributed/synced_batchnorm/test_groups.py
+++ b/tests/distributed/synced_batchnorm/test_groups.py
@@ -1,10 +1,12 @@
-import torch
+import argparse
+import os
+
 import numpy as np
-import apex
 import syncbn
-import os
-import argparse
-import torch.optim as optim
+import torch
+from torch import optim
+
+import apex
 
 
 def compare(desc, inp1, inp2, error):
diff --git a/tests/distributed/synced_batchnorm/two_gpu_test_different_batch_size.py b/tests/distributed/synced_batchnorm/two_gpu_test_different_batch_size.py
index a47c0f576..61fab27b0 100755
--- a/tests/distributed/synced_batchnorm/two_gpu_test_different_batch_size.py
+++ b/tests/distributed/synced_batchnorm/two_gpu_test_different_batch_size.py
@@ -1,11 +1,11 @@
-import torch
-import torch.nn as nn
-from torch.nn.parallel import DistributedDataParallel as DDP
-from apex.parallel import SyncBatchNorm as ApexSyncBatchNorm
-
 import argparse
 import os
+
 import numpy as np
+import torch
+from apex.parallel import SyncBatchNorm as ApexSyncBatchNorm
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
 
 var_batch = 16
 
@@ -33,7 +33,7 @@ def compare(desc, inp1, inp2, error=1e-5):
 torch.manual_seed(2809)
 # Setup DDP
 torch.cuda.set_device(args.local_rank)
-device = torch.device("cuda:{}".format(args.local_rank))
+device = torch.device(f"cuda:{args.local_rank}")
 
 torch.distributed.init_process_group(
     "nccl",
diff --git a/tests/distributed/synced_batchnorm/two_gpu_unit_test.py b/tests/distributed/synced_batchnorm/two_gpu_unit_test.py
index 3c97e9ac6..5c3534eb2 100644
--- a/tests/distributed/synced_batchnorm/two_gpu_unit_test.py
+++ b/tests/distributed/synced_batchnorm/two_gpu_unit_test.py
@@ -1,10 +1,12 @@
-import torch
+import argparse
+import os
+
 import numpy as np
-import apex
 import syncbn
-import os
-import argparse
-import torch.optim as optim
+import torch
+from torch import optim
+
+import apex
 
 
 def compare(desc, inp1, inp2, error):
@@ -119,7 +121,7 @@ def compare(desc, inp1, inp2, error):
 
 count = [
     space_size**2 * ((i + 1) * batch_size // args.world_size - i * batch_size // args.world_size)
-    for i in range(0, args.world_size)
+    for i in range(args.world_size)
 ]
 count = torch.cuda.IntTensor(count)