formatting

mlcommons · Oct 17, 2024 · 16484db · 16484db
1 parent 6cb1e05
commit 16484db
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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
+      module.momentum = module.momentum_backup