diff --git a/algorithmic_efficiency/pytorch_utils.py b/algorithmic_efficiency/pytorch_utils.py
index 2e5828912..590f500fa 100644
--- a/algorithmic_efficiency/pytorch_utils.py
+++ b/algorithmic_efficiency/pytorch_utils.py
@@ -57,6 +57,7 @@ def sync_ddp_time(time: float, device: torch.device) -> float:
   dist.all_reduce(time_tensor, op=dist.ReduceOp.MAX)
   return time_tensor.item()
 
+
 def update_batch_norm_fn(module: spec.ParameterContainer,
                          update_batch_norm: bool) -> None:
   bn_layers = (
@@ -75,4 +76,4 @@ def update_batch_norm_fn(module: spec.ParameterContainer,
       else:
         module.momentum = 0.0
     elif hasattr(module, 'momentum_backup'):
-      module.momentum = module.momentum_backup
\ No newline at end of file
+      module.momentum = module.momentum_backup