From a44253dd41fd5936117066b07c8a205875ae601e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 22 Mar 2024 21:23:04 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- nemo/collections/nlp/parts/nlp_overrides.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/nemo/collections/nlp/parts/nlp_overrides.py b/nemo/collections/nlp/parts/nlp_overrides.py index f0e9a40be1d0..a84785bcf407 100644 --- a/nemo/collections/nlp/parts/nlp_overrides.py +++ b/nemo/collections/nlp/parts/nlp_overrides.py @@ -442,9 +442,7 @@ def load_checkpoint(self, checkpoint_path: Union[str, Path]) -> Dict[str, Any]: # after dist_checkpointing.load, sharded tensors will be replaced with tensors checkpoint['state_dict'] = sharded_state_dict checkpoint['optimizer_states'] = [self.optimizer_sharded_state_dict()] - strategy = tensorstore.TensorStoreLoadShardedStrategy( - load_directly_on_device=True - ) + strategy = tensorstore.TensorStoreLoadShardedStrategy(load_directly_on_device=True) checkpoint = dist_checkpointing.load( sharded_state_dict=checkpoint, checkpoint_dir=checkpoint_path, sharded_strategy=strategy )