Apply isort and black reformatting

Signed-off-by: akoumpa <[email protected]>
NVIDIA · May 16, 2024 · 00f478e · 00f478e
1 parent d97fbcd
commit 00f478e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nemo/core/optim/mcore_optim.py b/nemo/core/optim/mcore_optim.py
@@ -56,7 +56,9 @@ def load_state_dict(self, state_dict):
  self.mcore_optimizer.load_state_dict(state_dict)
 
  def sharded_state_dict(self, model_sharded_state_dict, optimizer_state_dict=None, is_loading=False, **kwargs):
- sharding_type = 'fully_sharded_bucket_space' if kwargs.get('dist_ckpt_parallel_save', False) else 'dp_zero_gather_scatter'
+ sharding_type = (
+ 'fully_sharded_bucket_space' if kwargs.get('dist_ckpt_parallel_save', False) else 'dp_zero_gather_scatter'
+ )
  return self.mcore_optimizer.sharded_state_dict(
  model_sharded_state_dict, is_loading=False, sharding_type=sharding_type
  )