Update setup logic in training type plugins (data-parallel) [3 / n] (#10010)

Co-authored-by: thomas chaton <[email protected]>

Author: awaelchli

CHANGELOG.md

@@ -203,6 +203,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
     * Added `PrecisionPlugin.forward_context`, making it the default implementation for all `{train,val,test,predict}_step_context()` methods ([#9988](https://github.com/PyTorchLightning/pytorch-lightning/pull/9988))
     * Added `DDPSpawnPlugin.spawn()` for spawning new processes of a given function ([#10018](https://github.com/PyTorchLightning/pytorch-lightning/pull/10018))
     * Added `TrainingTypePlugin.{_setup_model, _setup_optimizer}` methods ([#9994](https://github.com/PyTorchLightning/pytorch-lightning/pull/9994))
+    * Implemented `DataParallelPlugin._setup_model` ([#10010](https://github.com/PyTorchLightning/pytorch-lightning/pull/10010))
 
 ### Changed
 

pytorch_lightning/plugins/training_type/dp.py

@@ -14,7 +14,7 @@
 from typing import List, Optional
 
 import torch
-from torch.nn import DataParallel
+from torch.nn import DataParallel, Module
 
 from pytorch_lightning.overrides.data_parallel import LightningParallelModule
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
@@ -54,7 +54,11 @@ def world_size(self) -> int:
     def setup(self) -> None:
         # model needs to be moved to the device before it is wrapped
         self.model_to_device()
-        self._model = DataParallel(LightningParallelModule(self._model), self.parallel_devices)
+        self._model = self._setup_model(LightningParallelModule(self._model))
+
+    def _setup_model(self, model: Module) -> DataParallel:
+        """Wraps the given model into a :class:`~torch.nn.parallel.DataParallel` module."""
+        return DataParallel(module=model, device_ids=self.parallel_devices)
 
     def reduce(self, collection: _METRIC_COLLECTION, *args, **kwargs) -> _METRIC_COLLECTION:
         """Reduces a collection of tensors from all processes. It can be applied to just a single tensor.