Address code review

Author: SeanNaren

pytorch_lightning/accelerators/accelerator.py

@@ -420,11 +420,11 @@ def optimizer_state(self, optimizer: Optimizer) -> Dict[str, Tensor]:
         """
         return getattr(self.training_type_plugin, 'optimizer_state', lambda x: x.state_dict())(optimizer)
 
-    def state_dict(self) -> Dict[str, Union[Any, Tensor]]:
+    def lightning_module_state_dict(self) -> Dict[str, Union[Any, Tensor]]:
         """
         Returns state of model. Allows for syncing/collating model state from processes in custom plugins.
         """
-        return self.training_type_plugin.state_dict()
+        return self.training_type_plugin.lightning_module_state_dict()
 
     def on_save(self, checkpoint: Dict[str, Union[Any, Tensor]]) -> Dict[str, Union[Any, Tensor]]:
         return self.training_type_plugin.on_save(checkpoint)

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -242,7 +242,7 @@ def update_global_step(self, total_batch_idx: int, current_global_step: int) ->
         """
         return current_global_step + 1
 
-    def state_dict(self) -> Dict[str, Union[Any, Tensor]]:
+    def lightning_module_state_dict(self) -> Dict[str, Union[Any, Tensor]]:
         """Returns model state."""
         model = self.lightning_module
         return model.state_dict()

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -273,7 +273,7 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
             'epoch': current_epoch,
             'global_step': global_step,
             'pytorch-lightning_version': pytorch_lightning.__version__,
-            'state_dict': self.trainer.accelerator.state_dict(),
+            'state_dict': self.trainer.accelerator.lightning_module_state_dict(),
         }
 
         if not weights_only: