Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Mar 22, 2024
1 parent 1650e59 commit a44253d
Showing 1 changed file with 1 addition and 3 deletions.
4 changes: 1 addition & 3 deletions nemo/collections/nlp/parts/nlp_overrides.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down

0 comments on commit a44253d

Please sign in to comment.