Address code review

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: