Update fully_sharded.py

Author: ananthsub

pytorch_lightning/strategies/fully_sharded.py

@@ -135,16 +135,15 @@ def setup_distributed(self) -> None:
 
     def setup(self, trainer: "pl.Trainer") -> None:
         self.accelerator.setup(trainer)
-        self.setup_optimizers(trainer)
-        self.setup_precision_plugin()
-        optimizers_to_device(self.optimizers, self.root_device)
 
         if self._layer_sync:
             self.model = self._layer_sync.apply(self.model)
 
         self.configure_ddp()
         self.barrier()
         self.setup_optimizers(trainer)
+        optimizers_to_device(self.optimizers, self.root_device)
+        self.setup_precision_plugin()
 
     @contextlib.contextmanager
     def model_sharded_context(self) -> Generator:
@@ -181,9 +180,6 @@ def configure_ddp(self) -> None:
             # (TODO: need to figure out solution)
             self.model_to_device()
 
-        # setup optimizers after fully sharded has wrapped the lightning module
-        self.setup_optimizers(self.lightning_module.trainer)
-
     def model_to_device(self) -> None:
         log.detail(f"{self.__class__.__name__}: moving model to device [{self.root_device}]...")
         # ensure we update the device type in the lightning module