pull from adrian's commit

Author: four4fish

pytorch_lightning/accelerators/cpu.py

@@ -34,9 +34,7 @@ def setup(self, trainer: "pl.Trainer") -> None:
                 If the selected device is not CPU.
         """
         if "cpu" not in str(self.root_device):
-            raise MisconfigurationException(
-                f"Device should be CPU, got {self.root_device} instead."
-            )
+            raise MisconfigurationException(f"Device should be CPU, got {self.root_device} instead.")
 
         return super().setup(trainer)
 

pytorch_lightning/accelerators/gpu.py

@@ -38,9 +38,7 @@ def setup_environment(self) -> None:
         """
         super().setup_environment()
         if "cuda" not in str(self.root_device):
-            raise MisconfigurationException(
-                f"Device should be GPU, got {self.root_device} instead"
-            )
+            raise MisconfigurationException(f"Device should be GPU, got {self.root_device} instead")
         torch.cuda.set_device(self.root_device)
 
     def setup(self, trainer: "pl.Trainer") -> None:

pytorch_lightning/plugins/training_type/single_tpu.py

@@ -42,7 +42,9 @@ def __init__(
 
         device = xm.xla_device(device)
         checkpoint_io = checkpoint_io or XLACheckpointIO()
-        super().__init__(accelerator=accelerator, device=device, checkpoint_io=checkpoint_io, precision_plugin=precision_plugin)
+        super().__init__(
+            accelerator=accelerator, device=device, checkpoint_io=checkpoint_io, precision_plugin=precision_plugin
+        )
 
         self.debug = debug
         self.tpu_local_core_rank = 0

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -64,7 +64,9 @@ def __init__(
         checkpoint_io = checkpoint_io or XLACheckpointIO()
         super().__init__(
             accelerator=accelerator,
-            parallel_devices=parallel_devices, checkpoint_io=checkpoint_io, precision_plugin=precision_plugin,
+            parallel_devices=parallel_devices,
+            checkpoint_io=checkpoint_io,
+            precision_plugin=precision_plugin,
         )
         self.debug = debug
         self.tpu_local_core_rank = 0

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -42,7 +42,10 @@ class TrainingTypePlugin(ABC):
     loop."""
 
     def __init__(
-        self, accelerator: Optional["pl.Accelerator"] = None, checkpoint_io: Optional[CheckpointIO] = None, precision_plugin: Optional[PrecisionPlugin] = None
+        self,
+        accelerator: Optional["pl.Accelerator"] = None,
+        checkpoint_io: Optional[CheckpointIO] = None,
+        precision_plugin: Optional[PrecisionPlugin] = None,
     ) -> None:
         self._accelerator = accelerator
         self._model: Optional[Module] = None

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -691,10 +691,16 @@ def select_precision_plugin(self) -> PrecisionPlugin:
 
     def create_training_type_plugin(self) -> TrainingTypePlugin:
         if self.use_ddp2:
-            plugin = DDP2Plugin(accelerator=self.accelerator, parallel_devices=self.parallel_devices, cluster_environment=self.cluster_environment,)
+            plugin = DDP2Plugin(
+                accelerator=self.accelerator,
+                parallel_devices=self.parallel_devices,
+                cluster_environment=self.cluster_environment,
+            )
         elif self.use_ddp and self.use_deepspeed:
             plugin = DeepSpeedPlugin(
-                accelerator=self.accelerator, cluster_environment=self.select_cluster_environment(), parallel_devices=self.parallel_devices,
+                accelerator=self.accelerator,
+                cluster_environment=self.select_cluster_environment(),
+                parallel_devices=self.parallel_devices,
             )
         elif self.use_ddp:
             use_slurm_ddp = self.use_ddp and self._is_slurm_managing_tasks()
@@ -733,7 +739,9 @@ def create_training_type_plugin(self) -> TrainingTypePlugin:
                 ddp_plugin_cls = DDPPlugin
 
             plugin = ddp_plugin_cls(
-                accelerator=self.accelerator, parallel_devices=self.parallel_devices, cluster_environment=self.cluster_environment,
+                accelerator=self.accelerator,
+                parallel_devices=self.parallel_devices,
+                cluster_environment=self.cluster_environment,
             )
         elif self.use_dp:
             plugin = DataParallelPlugin(accelerator=self.accelerator, parallel_devices=self.parallel_devices)
@@ -745,7 +753,10 @@ def create_training_type_plugin(self) -> TrainingTypePlugin:
             plugin = IPUPlugin(accelerator=self.accelerator, parallel_devices=self.parallel_devices)
         else:
             single_gpu_ordinal = device_parser.determine_root_gpu_device(self.parallel_device_ids)
-            plugin = SingleDevicePlugin(accelerator=self.accelerator, device=(torch.device(f"cuda:{single_gpu_ordinal}" if self.use_gpu else "cpu")),)
+            plugin = SingleDevicePlugin(
+                accelerator=self.accelerator,
+                device=(torch.device(f"cuda:{single_gpu_ordinal}" if self.use_gpu else "cpu")),
+            )
         return plugin
 
     def resolve_training_type_plugin(self, training_type: TrainingTypePlugin) -> TrainingTypePlugin:

tests/accelerators/test_cpu.py

@@ -16,7 +16,9 @@
 
 def test_restore_checkpoint_after_pre_dispatch_default():
     """Assert default for restore_checkpoint_after_pre_dispatch is False."""
-    plugin = SingleDevicePlugin(accelerator=CPUAccelerator(), device=torch.device("cpu"), precision_plugin=PrecisionPlugin())
+    plugin = SingleDevicePlugin(
+        accelerator=CPUAccelerator(), device=torch.device("cpu"), precision_plugin=PrecisionPlugin()
+    )
     assert not plugin.restore_checkpoint_after_pre_dispatch
     assert not plugin.restore_checkpoint_after_pre_dispatch
 
@@ -48,7 +50,12 @@ def load_checkpoint(self, checkpoint_path: Union[str, Path]) -> Dict[str, Any]:
     checkpoint_path = os.path.join(tmpdir, "model.pt")
     trainer.save_checkpoint(checkpoint_path)
 
-    plugin = TestPlugin(accelerator=CPUAccelerator(), precision_plugin=PrecisionPlugin(), device=torch.device("cpu"), checkpoint_io=TorchCheckpointIO())
+    plugin = TestPlugin(
+        accelerator=CPUAccelerator(),
+        precision_plugin=PrecisionPlugin(),
+        device=torch.device("cpu"),
+        checkpoint_io=TorchCheckpointIO(),
+    )
     assert plugin.restore_checkpoint_after_pre_dispatch == restore_after_pre_dispatch
 
     trainer = Trainer(default_root_dir=tmpdir, strategy=plugin, fast_dev_run=True)

tests/accelerators/test_tpu.py

@@ -13,7 +13,7 @@
 # limitations under the License
 import collections
 from copy import deepcopy
-from unittest.mock import patch, Mock
+from unittest.mock import Mock, patch
 
 import pytest
 import torch