From 16484dbd46fa74e569089ba51f99eeef4746dcc4 Mon Sep 17 00:00:00 2001 From: Priya Kasimbeg Date: Thu, 17 Oct 2024 18:31:14 +0000 Subject: [PATCH] formatting --- algorithmic_efficiency/pytorch_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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