address comment

huggingface · Nov 21, 2023 · d18b7e3 · d18b7e3
1 parent cbb25ea
commit d18b7e3
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -2986,13 +2986,8 @@ def get_state_dict(self, model, unwrap=True):
 
                 state_dict = clone_tensors_for_torch_save(self.unwrap_model(model).state_dict())
         elif self.distributed_type == DistributedType.FSDP:
-            from torch.distributed.fsdp import (
-                FullStateDictConfig,
-                StateDictType,
-            )
-            from torch.distributed.fsdp import (
-                FullyShardedDataParallel as FSDP,
-            )
+            from torch.distributed.fsdp import FullStateDictConfig, StateDictType
+            from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
             full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
             with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, full_state_dict_config):