From 0d1f76dd819347cddb63223eb32a9266b2c1486c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Sat, 13 Feb 2021 01:23:32 +0100
Subject: [PATCH 01/12] flake8

---
 pytorch_lightning/accelerators/accelerator.py |  5 ++--
 pytorch_lightning/plugins/base_plugin.py      |  3 +-
 .../plugins/precision/precision_plugin.py     |  3 +-
 .../plugins/precision/tpu_bfloat.py           |  2 +-
 .../plugins/training_type/__init__.py         | 28 +++++++++----------
 .../plugins/training_type/ddp.py              |  3 +-
 .../plugins/training_type/horovod.py          |  2 +-
 .../plugins/training_type/parallel.py         |  1 -
 .../plugins/training_type/rpc.py              |  1 -
 .../plugins/training_type/rpc_sequential.py   |  1 +
 .../plugins/training_type/sharded_spawn.py    |  1 -
 .../plugins/training_type/single_device.py    |  1 -
 .../plugins/training_type/single_tpu.py       |  4 +--
 .../plugins/training_type/tpu_spawn.py        |  4 +--
 .../training_type/training_type_plugin.py     |  1 -
 setup.cfg                                     |  5 ----
 16 files changed, 26 insertions(+), 39 deletions(-)

diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index e348a57b5c103..77062f350ca09 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -11,21 +11,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Iterable, Optional, TYPE_CHECKING, Union
+from typing import Any, Callable, Iterable, Optional, Union
 
 import torch
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
 
 from pytorch_lightning.core import LightningModule
+from pytorch_lightning.trainer.trainer import Trainer
 from pytorch_lightning.plugins.precision import (
     ApexMixedPrecisionPlugin,
-    MixedPrecisionPlugin,
     NativeMixedPrecisionPlugin,
     PrecisionPlugin,
 )
 from pytorch_lightning.plugins.training_type import TrainingTypePlugin
-from pytorch_lightning.plugins.training_type.horovod import HorovodPlugin
 from pytorch_lightning.utilities.apply_func import move_data_to_device
 from pytorch_lightning.utilities.distributed import all_gather_ddp_if_available
 from pytorch_lightning.utilities.enums import AMPType, LightningEnum
diff --git a/pytorch_lightning/plugins/base_plugin.py b/pytorch_lightning/plugins/base_plugin.py
index b8bdf38a57137..3dd803d4b6571 100644
--- a/pytorch_lightning/plugins/base_plugin.py
+++ b/pytorch_lightning/plugins/base_plugin.py
@@ -13,9 +13,8 @@
 # limitations under the License.
 import contextlib
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Generator, Optional, overload, Sequence, Tuple
+from typing import Generator, Optional, Sequence, Tuple
 
-import torch
 from torch.nn import Module
 
 
diff --git a/pytorch_lightning/plugins/precision/precision_plugin.py b/pytorch_lightning/plugins/precision/precision_plugin.py
index 2216d3ae46d53..34879e514a6f2 100644
--- a/pytorch_lightning/plugins/precision/precision_plugin.py
+++ b/pytorch_lightning/plugins/precision/precision_plugin.py
@@ -24,7 +24,8 @@
 
 class PrecisionPlugin(Plugin):
     """ Plugin handling the precision-specific parts of the training.
-    The static classattributes EPSILON and precision must be overwritten in child-classes and their default values reflect fp32 training
+    The static classattributes EPSILON and precision must be overwritten in child-classes and their
+    default values reflect fp32 training.
     """
     EPSILON = 1e-6
     precision = 32
diff --git a/pytorch_lightning/plugins/precision/tpu_bfloat.py b/pytorch_lightning/plugins/precision/tpu_bfloat.py
index c911bf69184f6..7f4916dd26a46 100644
--- a/pytorch_lightning/plugins/precision/tpu_bfloat.py
+++ b/pytorch_lightning/plugins/precision/tpu_bfloat.py
@@ -25,4 +25,4 @@ class TPUHalfPrecisionPlugin(PrecisionPlugin):
 
     def connect(self, model: torch.nn.Module, optimizers, lr_schedulers):
         os.environ["XLA_USE_BF16"] = str(1)
-        return super().connect(model=model, optimizers=optimizers, lr_schedulers=lr_schedulers)
\ No newline at end of file
+        return super().connect(model=model, optimizers=optimizers, lr_schedulers=lr_schedulers)
diff --git a/pytorch_lightning/plugins/training_type/__init__.py b/pytorch_lightning/plugins/training_type/__init__.py
index a5a644fc6568c..c50e42d48b953 100644
--- a/pytorch_lightning/plugins/training_type/__init__.py
+++ b/pytorch_lightning/plugins/training_type/__init__.py
@@ -1,14 +1,14 @@
-from pytorch_lightning.plugins.training_type.ddp import DDPPlugin
-from pytorch_lightning.plugins.training_type.ddp2 import DDP2Plugin
-from pytorch_lightning.plugins.training_type.ddp_spawn import DDPSpawnPlugin
-from pytorch_lightning.plugins.training_type.dp import DataParallelPlugin
-from pytorch_lightning.plugins.training_type.horovod import HorovodPlugin
-from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin
-from pytorch_lightning.plugins.training_type.rpc import RPCPlugin
-from pytorch_lightning.plugins.training_type.rpc_sequential import RPCSequentialPlugin
-from pytorch_lightning.plugins.training_type.sharded import DDPShardedPlugin
-from pytorch_lightning.plugins.training_type.sharded_spawn import DDPSpawnShardedPlugin
-from pytorch_lightning.plugins.training_type.single_device import SingleDevicePlugin
-from pytorch_lightning.plugins.training_type.single_tpu import SingleTPUPlugin
-from pytorch_lightning.plugins.training_type.tpu_spawn import TPUSpawnPlugin
-from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
+from pytorch_lightning.plugins.training_type.ddp import DDPPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.ddp2 import DDP2Plugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.ddp_spawn import DDPSpawnPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.dp import DataParallelPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.horovod import HorovodPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.rpc import RPCPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.rpc_sequential import RPCSequentialPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.sharded import DDPShardedPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.sharded_spawn import DDPSpawnShardedPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.single_device import SingleDevicePlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.single_tpu import SingleTPUPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.tpu_spawn import TPUSpawnPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin  # noqa: F401
diff --git a/pytorch_lightning/plugins/training_type/ddp.py b/pytorch_lightning/plugins/training_type/ddp.py
index 52a24655f0846..fb4386fbc10f8 100644
--- a/pytorch_lightning/plugins/training_type/ddp.py
+++ b/pytorch_lightning/plugins/training_type/ddp.py
@@ -27,7 +27,6 @@
 from pytorch_lightning.distributed import LightningDistributed
 from pytorch_lightning.overrides import LightningDistributedModule
 from pytorch_lightning.overrides.distributed import prepare_for_backward
-from pytorch_lightning.plugins.environments import SLURMEnvironment, TorchElasticEnvironment
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin
 from pytorch_lightning.utilities import _HYDRA_AVAILABLE, _TORCH_GREATER_EQUAL_1_7, rank_zero_warn
@@ -120,7 +119,7 @@ def _call_children_scripts(self):
         command = sys.argv
         try:
             full_path = path_lib(command[0])
-        except Exception as e:
+        except Exception:
             full_path = os.path.abspath(command[0])
 
         command[0] = full_path
diff --git a/pytorch_lightning/plugins/training_type/horovod.py b/pytorch_lightning/plugins/training_type/horovod.py
index 2393c040bcc8f..c90e6ac5aff28 100644
--- a/pytorch_lightning/plugins/training_type/horovod.py
+++ b/pytorch_lightning/plugins/training_type/horovod.py
@@ -101,7 +101,7 @@ def start_training(self, trainer):
         hvd.join()
 
     def start_testing(self, trainer):
-        with ExitStack() as stack:
+        with ExitStack():
             # set up training routine
             # self.trainer.train_loop.setup_training(self.trainer.model)
             self._results = trainer.run_test()
diff --git a/pytorch_lightning/plugins/training_type/parallel.py b/pytorch_lightning/plugins/training_type/parallel.py
index a67dee93a6500..3d6d661f171ee 100644
--- a/pytorch_lightning/plugins/training_type/parallel.py
+++ b/pytorch_lightning/plugins/training_type/parallel.py
@@ -21,7 +21,6 @@
 
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.overrides.base import unwrap_lightning_module
-from pytorch_lightning.overrides.data_parallel import LightningDistributedDataParallel
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
 from pytorch_lightning.utilities.distributed import all_gather_ddp_if_available, ReduceOp
diff --git a/pytorch_lightning/plugins/training_type/rpc.py b/pytorch_lightning/plugins/training_type/rpc.py
index be81cd2a03c56..e2d4b4e081134 100644
--- a/pytorch_lightning/plugins/training_type/rpc.py
+++ b/pytorch_lightning/plugins/training_type/rpc.py
@@ -17,7 +17,6 @@
 
 import torch
 
-from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.ddp import DDPPlugin
 from pytorch_lightning.utilities import _RPC_AVAILABLE
diff --git a/pytorch_lightning/plugins/training_type/rpc_sequential.py b/pytorch_lightning/plugins/training_type/rpc_sequential.py
index 331cbe76639f3..ddee076c9ef0e 100644
--- a/pytorch_lightning/plugins/training_type/rpc_sequential.py
+++ b/pytorch_lightning/plugins/training_type/rpc_sequential.py
@@ -21,6 +21,7 @@
 from torch.nn.parallel import DistributedDataParallel
 from torch.optim import Optimizer
 
+from pytorch_lightning.trainer.trainer import Trainer
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.overrides.distributed import LightningDistributedModule
 from pytorch_lightning.plugins.training_type.rpc import DEFAULT_RPC_TIMEOUT_SEC, RPCPlugin
diff --git a/pytorch_lightning/plugins/training_type/sharded_spawn.py b/pytorch_lightning/plugins/training_type/sharded_spawn.py
index c38690473b77d..a8d497cd119b0 100644
--- a/pytorch_lightning/plugins/training_type/sharded_spawn.py
+++ b/pytorch_lightning/plugins/training_type/sharded_spawn.py
@@ -1,7 +1,6 @@
 from typing import Optional
 
 from pytorch_lightning.core.lightning import LightningModule
-from pytorch_lightning.core.optimizer import is_lightning_optimizer
 from pytorch_lightning.plugins.training_type.ddp_spawn import DDPSpawnPlugin
 from pytorch_lightning.utilities import _FAIRSCALE_AVAILABLE, rank_zero_only
 
diff --git a/pytorch_lightning/plugins/training_type/single_device.py b/pytorch_lightning/plugins/training_type/single_device.py
index 0eda31833d6fa..49ab0b4160692 100644
--- a/pytorch_lightning/plugins/training_type/single_device.py
+++ b/pytorch_lightning/plugins/training_type/single_device.py
@@ -1,7 +1,6 @@
 from typing import Any, Union
 
 import torch
-from torch._C import device
 
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
 
diff --git a/pytorch_lightning/plugins/training_type/single_tpu.py b/pytorch_lightning/plugins/training_type/single_tpu.py
index 46df404bdc02f..f4ae9ce22dbae 100644
--- a/pytorch_lightning/plugins/training_type/single_tpu.py
+++ b/pytorch_lightning/plugins/training_type/single_tpu.py
@@ -1,4 +1,3 @@
-import io
 import os
 from typing import Optional, Union
 
@@ -11,7 +10,6 @@
 from pytorch_lightning.utilities.apply_func import move_data_to_device
 
 if _TPU_AVAILABLE:
-    import torch_xla
     import torch_xla.core.xla_model as xm
 
 
@@ -68,4 +66,4 @@ def on_save(self, checkpoint: dict) -> dict:
 
     @property
     def is_distributed(self):
-        return False
\ No newline at end of file
+        return False
diff --git a/pytorch_lightning/plugins/training_type/tpu_spawn.py b/pytorch_lightning/plugins/training_type/tpu_spawn.py
index 4c5844da94ced..262330c84c15a 100644
--- a/pytorch_lightning/plugins/training_type/tpu_spawn.py
+++ b/pytorch_lightning/plugins/training_type/tpu_spawn.py
@@ -1,7 +1,7 @@
 import io
 import os
 import re
-from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Iterable, Optional, Sequence, Union
 
 import torch
 import torch.multiprocessing as mp
@@ -189,7 +189,7 @@ def post_training(self) -> None:
         # restore main state with best weights
         best_path = self.mp_queue.get()
         last_path = self.mp_queue.get()
-        results = self.mp_queue.get()
+        self._results = self.mp_queue.get()
 
         # transfer back the best path to the trainer
         if self.lightning_module.trainer.checkpoint_callback is not None:
diff --git a/pytorch_lightning/plugins/training_type/training_type_plugin.py b/pytorch_lightning/plugins/training_type/training_type_plugin.py
index db0e390c4b03e..ee25a21ba2d5b 100644
--- a/pytorch_lightning/plugins/training_type/training_type_plugin.py
+++ b/pytorch_lightning/plugins/training_type/training_type_plugin.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import os
 from abc import ABC, abstractmethod
 from typing import Any, Optional, TYPE_CHECKING, Union
 
diff --git a/setup.cfg b/setup.cfg
index 8b02b462fad84..338e2871e242d 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -67,11 +67,6 @@ exclude =
     *.egg
     build
     temp
-    # TODO: temporary until accelerator refactor finished
-    pytorch_lightning/accelerators/accelerator.py
-    pytorch_lightning/plugins/training_type
-    pytorch_lightning/plugins/precision
-    pytorch_lightning/plugins/base_plugin.py
 
 select = E,W,F
 doctests = True

From 5475f85c0be23756854e202e100c9ae324681482 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Sun, 14 Feb 2021 03:36:25 +0100
Subject: [PATCH 02/12] fix cyclic import

---
 pytorch_lightning/accelerators/accelerator.py             | 5 ++---
 pytorch_lightning/plugins/training_type/rpc_sequential.py | 5 ++---
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index 77062f350ca09..321c2fd78aefa 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -18,7 +18,6 @@
 from torch.utils.data import DataLoader
 
 from pytorch_lightning.core import LightningModule
-from pytorch_lightning.trainer.trainer import Trainer
 from pytorch_lightning.plugins.precision import (
     ApexMixedPrecisionPlugin,
     NativeMixedPrecisionPlugin,
@@ -63,7 +62,7 @@ def __init__(
         self.lr_schedulers = None
         self.optimizer_frequencies = None
 
-    def setup(self, trainer: "Trainer", model: LightningModule) -> None:
+    def setup(self, trainer, model: LightningModule) -> None:
         """
         Connects the plugins to the training process, creates optimizers
 
@@ -302,7 +301,7 @@ def on_train_end(self) -> None:
         """Hook to do something at the end of the training"""
         pass
 
-    def setup_optimizers(self, trainer: "Trainer"):
+    def setup_optimizers(self, trainer):
         """creates optimizers and schedulers
 
         Args:
diff --git a/pytorch_lightning/plugins/training_type/rpc_sequential.py b/pytorch_lightning/plugins/training_type/rpc_sequential.py
index ddee076c9ef0e..522475e64caf1 100644
--- a/pytorch_lightning/plugins/training_type/rpc_sequential.py
+++ b/pytorch_lightning/plugins/training_type/rpc_sequential.py
@@ -21,7 +21,6 @@
 from torch.nn.parallel import DistributedDataParallel
 from torch.optim import Optimizer
 
-from pytorch_lightning.trainer.trainer import Trainer
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.overrides.distributed import LightningDistributedModule
 from pytorch_lightning.plugins.training_type.rpc import DEFAULT_RPC_TIMEOUT_SEC, RPCPlugin
@@ -330,11 +329,11 @@ def post_training(self):
         if self.main_rpc_process:
             super().post_training()
 
-    def start_training(self, trainer: 'Trainer') -> None:
+    def start_training(self, trainer) -> None:
         if self.main_rpc_process:
             super().start_training(trainer)
 
-    def start_testing(self, trainer: 'Trainer') -> None:
+    def start_testing(self, trainer) -> None:
         if self.main_rpc_process:
             super().start_testing(trainer)
 

From 36afc887fc4b2c706540b428d31c0883f7d3aeac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Sun, 14 Feb 2021 03:38:11 +0100
Subject: [PATCH 03/12] isort

---
 pytorch_lightning/accelerators/accelerator.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index 321c2fd78aefa..e2a9ca6492af9 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -18,11 +18,7 @@
 from torch.utils.data import DataLoader
 
 from pytorch_lightning.core import LightningModule
-from pytorch_lightning.plugins.precision import (
-    ApexMixedPrecisionPlugin,
-    NativeMixedPrecisionPlugin,
-    PrecisionPlugin,
-)
+from pytorch_lightning.plugins.precision import ApexMixedPrecisionPlugin, NativeMixedPrecisionPlugin, PrecisionPlugin
 from pytorch_lightning.plugins.training_type import TrainingTypePlugin
 from pytorch_lightning.utilities.apply_func import move_data_to_device
 from pytorch_lightning.utilities.distributed import all_gather_ddp_if_available

From c18335bcaf3bb229bd78874ad0f78bd08723c7eb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Sat, 13 Feb 2021 01:23:32 +0100
Subject: [PATCH 04/12] flake8

---
 pytorch_lightning/accelerators/accelerator.py |  3 +-
 pytorch_lightning/plugins/base_plugin.py      |  3 +-
 .../plugins/precision/precision_plugin.py     |  3 +-
 .../plugins/precision/tpu_bfloat.py           |  2 +-
 .../plugins/training_type/__init__.py         | 28 +++++++++----------
 .../plugins/training_type/ddp.py              |  3 +-
 .../plugins/training_type/horovod.py          |  4 +--
 .../plugins/training_type/rpc.py              |  1 -
 .../plugins/training_type/rpc_sequential.py   |  1 +
 .../plugins/training_type/sharded_spawn.py    |  1 -
 .../plugins/training_type/single_device.py    |  1 -
 .../plugins/training_type/single_tpu.py       |  4 +--
 .../plugins/training_type/tpu_spawn.py        |  4 +--
 setup.cfg                                     |  5 ----
 14 files changed, 26 insertions(+), 37 deletions(-)

diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index 967b6a85c878b..fcd0f9229acc7 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -18,14 +18,13 @@
 from torch.utils.data import DataLoader
 
 from pytorch_lightning.core import LightningModule
+from pytorch_lightning.trainer.trainer import Trainer
 from pytorch_lightning.plugins.precision import (
     ApexMixedPrecisionPlugin,
-    MixedPrecisionPlugin,
     NativeMixedPrecisionPlugin,
     PrecisionPlugin,
 )
 from pytorch_lightning.plugins.training_type import TrainingTypePlugin
-from pytorch_lightning.plugins.training_type.horovod import HorovodPlugin
 from pytorch_lightning.utilities.apply_func import move_data_to_device
 from pytorch_lightning.utilities.distributed import all_gather_ddp_if_available
 from pytorch_lightning.utilities.enums import AMPType, LightningEnum
diff --git a/pytorch_lightning/plugins/base_plugin.py b/pytorch_lightning/plugins/base_plugin.py
index e495d9ffadc3f..19033a20997df 100644
--- a/pytorch_lightning/plugins/base_plugin.py
+++ b/pytorch_lightning/plugins/base_plugin.py
@@ -13,9 +13,8 @@
 # limitations under the License.
 import contextlib
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Generator, Optional, overload, Sequence, Tuple
+from typing import Generator, Optional, Sequence, Tuple
 
-import torch
 from torch.nn import Module
 
 
diff --git a/pytorch_lightning/plugins/precision/precision_plugin.py b/pytorch_lightning/plugins/precision/precision_plugin.py
index 2216d3ae46d53..34879e514a6f2 100644
--- a/pytorch_lightning/plugins/precision/precision_plugin.py
+++ b/pytorch_lightning/plugins/precision/precision_plugin.py
@@ -24,7 +24,8 @@
 
 class PrecisionPlugin(Plugin):
     """ Plugin handling the precision-specific parts of the training.
-    The static classattributes EPSILON and precision must be overwritten in child-classes and their default values reflect fp32 training
+    The static classattributes EPSILON and precision must be overwritten in child-classes and their
+    default values reflect fp32 training.
     """
     EPSILON = 1e-6
     precision = 32
diff --git a/pytorch_lightning/plugins/precision/tpu_bfloat.py b/pytorch_lightning/plugins/precision/tpu_bfloat.py
index c911bf69184f6..7f4916dd26a46 100644
--- a/pytorch_lightning/plugins/precision/tpu_bfloat.py
+++ b/pytorch_lightning/plugins/precision/tpu_bfloat.py
@@ -25,4 +25,4 @@ class TPUHalfPrecisionPlugin(PrecisionPlugin):
 
     def connect(self, model: torch.nn.Module, optimizers, lr_schedulers):
         os.environ["XLA_USE_BF16"] = str(1)
-        return super().connect(model=model, optimizers=optimizers, lr_schedulers=lr_schedulers)
\ No newline at end of file
+        return super().connect(model=model, optimizers=optimizers, lr_schedulers=lr_schedulers)
diff --git a/pytorch_lightning/plugins/training_type/__init__.py b/pytorch_lightning/plugins/training_type/__init__.py
index a5a644fc6568c..c50e42d48b953 100644
--- a/pytorch_lightning/plugins/training_type/__init__.py
+++ b/pytorch_lightning/plugins/training_type/__init__.py
@@ -1,14 +1,14 @@
-from pytorch_lightning.plugins.training_type.ddp import DDPPlugin
-from pytorch_lightning.plugins.training_type.ddp2 import DDP2Plugin
-from pytorch_lightning.plugins.training_type.ddp_spawn import DDPSpawnPlugin
-from pytorch_lightning.plugins.training_type.dp import DataParallelPlugin
-from pytorch_lightning.plugins.training_type.horovod import HorovodPlugin
-from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin
-from pytorch_lightning.plugins.training_type.rpc import RPCPlugin
-from pytorch_lightning.plugins.training_type.rpc_sequential import RPCSequentialPlugin
-from pytorch_lightning.plugins.training_type.sharded import DDPShardedPlugin
-from pytorch_lightning.plugins.training_type.sharded_spawn import DDPSpawnShardedPlugin
-from pytorch_lightning.plugins.training_type.single_device import SingleDevicePlugin
-from pytorch_lightning.plugins.training_type.single_tpu import SingleTPUPlugin
-from pytorch_lightning.plugins.training_type.tpu_spawn import TPUSpawnPlugin
-from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
+from pytorch_lightning.plugins.training_type.ddp import DDPPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.ddp2 import DDP2Plugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.ddp_spawn import DDPSpawnPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.dp import DataParallelPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.horovod import HorovodPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.rpc import RPCPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.rpc_sequential import RPCSequentialPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.sharded import DDPShardedPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.sharded_spawn import DDPSpawnShardedPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.single_device import SingleDevicePlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.single_tpu import SingleTPUPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.tpu_spawn import TPUSpawnPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin  # noqa: F401
diff --git a/pytorch_lightning/plugins/training_type/ddp.py b/pytorch_lightning/plugins/training_type/ddp.py
index 15b456f8e8391..f018b47533f4a 100644
--- a/pytorch_lightning/plugins/training_type/ddp.py
+++ b/pytorch_lightning/plugins/training_type/ddp.py
@@ -27,7 +27,6 @@
 from pytorch_lightning.distributed import LightningDistributed
 from pytorch_lightning.overrides import LightningDistributedModule
 from pytorch_lightning.overrides.distributed import prepare_for_backward
-from pytorch_lightning.plugins.environments import SLURMEnvironment, TorchElasticEnvironment
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin
 from pytorch_lightning.utilities import _HYDRA_AVAILABLE, _TORCH_GREATER_EQUAL_1_7, rank_zero_warn
@@ -120,7 +119,7 @@ def _call_children_scripts(self):
         command = sys.argv
         try:
             full_path = path_lib(command[0])
-        except Exception as e:
+        except Exception:
             full_path = os.path.abspath(command[0])
 
         command[0] = full_path
diff --git a/pytorch_lightning/plugins/training_type/horovod.py b/pytorch_lightning/plugins/training_type/horovod.py
index c1de2d7833177..e6281c48da75d 100644
--- a/pytorch_lightning/plugins/training_type/horovod.py
+++ b/pytorch_lightning/plugins/training_type/horovod.py
@@ -101,14 +101,14 @@ def start_training(self, trainer):
         hvd.join()
 
     def start_testing(self, trainer):
-        with ExitStack() as stack:
+        with ExitStack():
             self._results = trainer.run_test()
 
         # Make sure all workers have finished training before returning to the user
         hvd.join()
 
     def start_predicting(self, trainer):
-        with ExitStack() as stack:
+        with ExitStack():
             # set up training routine
             self._results = trainer.run_predict()
 
diff --git a/pytorch_lightning/plugins/training_type/rpc.py b/pytorch_lightning/plugins/training_type/rpc.py
index be81cd2a03c56..e2d4b4e081134 100644
--- a/pytorch_lightning/plugins/training_type/rpc.py
+++ b/pytorch_lightning/plugins/training_type/rpc.py
@@ -17,7 +17,6 @@
 
 import torch
 
-from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.ddp import DDPPlugin
 from pytorch_lightning.utilities import _RPC_AVAILABLE
diff --git a/pytorch_lightning/plugins/training_type/rpc_sequential.py b/pytorch_lightning/plugins/training_type/rpc_sequential.py
index fc707afb3e2c2..b9a352aedb387 100644
--- a/pytorch_lightning/plugins/training_type/rpc_sequential.py
+++ b/pytorch_lightning/plugins/training_type/rpc_sequential.py
@@ -21,6 +21,7 @@
 from torch.nn.parallel import DistributedDataParallel
 from torch.optim import Optimizer
 
+from pytorch_lightning.trainer.trainer import Trainer
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.overrides.distributed import LightningDistributedModule
 from pytorch_lightning.plugins.training_type.rpc import DEFAULT_RPC_TIMEOUT_SEC, RPCPlugin
diff --git a/pytorch_lightning/plugins/training_type/sharded_spawn.py b/pytorch_lightning/plugins/training_type/sharded_spawn.py
index c38690473b77d..a8d497cd119b0 100644
--- a/pytorch_lightning/plugins/training_type/sharded_spawn.py
+++ b/pytorch_lightning/plugins/training_type/sharded_spawn.py
@@ -1,7 +1,6 @@
 from typing import Optional
 
 from pytorch_lightning.core.lightning import LightningModule
-from pytorch_lightning.core.optimizer import is_lightning_optimizer
 from pytorch_lightning.plugins.training_type.ddp_spawn import DDPSpawnPlugin
 from pytorch_lightning.utilities import _FAIRSCALE_AVAILABLE, rank_zero_only
 
diff --git a/pytorch_lightning/plugins/training_type/single_device.py b/pytorch_lightning/plugins/training_type/single_device.py
index 0eda31833d6fa..49ab0b4160692 100644
--- a/pytorch_lightning/plugins/training_type/single_device.py
+++ b/pytorch_lightning/plugins/training_type/single_device.py
@@ -1,7 +1,6 @@
 from typing import Any, Union
 
 import torch
-from torch._C import device
 
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
 
diff --git a/pytorch_lightning/plugins/training_type/single_tpu.py b/pytorch_lightning/plugins/training_type/single_tpu.py
index 40fc9fba3a6a7..3ddfd98128787 100644
--- a/pytorch_lightning/plugins/training_type/single_tpu.py
+++ b/pytorch_lightning/plugins/training_type/single_tpu.py
@@ -1,4 +1,3 @@
-import io
 import os
 from typing import Optional, Union
 
@@ -11,7 +10,6 @@
 from pytorch_lightning.utilities.apply_func import move_data_to_device
 
 if _TPU_AVAILABLE:
-    import torch_xla
     import torch_xla.core.xla_model as xm
 
 
@@ -68,4 +66,4 @@ def on_save(self, checkpoint: dict) -> dict:
 
     @property
     def is_distributed(self):
-        return False
\ No newline at end of file
+        return False
diff --git a/pytorch_lightning/plugins/training_type/tpu_spawn.py b/pytorch_lightning/plugins/training_type/tpu_spawn.py
index 0136e78a4381f..83a441aabad30 100644
--- a/pytorch_lightning/plugins/training_type/tpu_spawn.py
+++ b/pytorch_lightning/plugins/training_type/tpu_spawn.py
@@ -1,7 +1,7 @@
 import io
 import os
 import re
-from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Iterable, Optional, Sequence, Union
 
 import torch
 import torch.multiprocessing as mp
@@ -201,7 +201,7 @@ def post_dispatch(self) -> None:
         # restore main state with best weights
         best_path = self.mp_queue.get()
         last_path = self.mp_queue.get()
-        results = self.mp_queue.get()
+        self._results = self.mp_queue.get()
 
         # transfer back the best path to the trainer
         if self.lightning_module.trainer.checkpoint_callback is not None:
diff --git a/setup.cfg b/setup.cfg
index f622581b5aaf7..be8f7cd50fa2d 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -67,11 +67,6 @@ exclude =
     *.egg
     build
     temp
-    # TODO: temporary until accelerator refactor finished
-    pytorch_lightning/accelerators/accelerator.py
-    pytorch_lightning/plugins/training_type
-    pytorch_lightning/plugins/precision
-    pytorch_lightning/plugins/base_plugin.py
 
 select = E,W,F
 doctests = True

From a6534250bd9670dd980d1bd1208e806a6abe8259 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Sun, 14 Feb 2021 03:36:25 +0100
Subject: [PATCH 05/12] fix cyclic import

---
 pytorch_lightning/accelerators/accelerator.py             | 5 ++---
 pytorch_lightning/plugins/training_type/rpc_sequential.py | 5 ++---
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index fcd0f9229acc7..2e876d21d2481 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -18,7 +18,6 @@
 from torch.utils.data import DataLoader
 
 from pytorch_lightning.core import LightningModule
-from pytorch_lightning.trainer.trainer import Trainer
 from pytorch_lightning.plugins.precision import (
     ApexMixedPrecisionPlugin,
     NativeMixedPrecisionPlugin,
@@ -63,7 +62,7 @@ def __init__(
         self.lr_schedulers = None
         self.optimizer_frequencies = None
 
-    def setup(self, trainer: "Trainer", model: LightningModule) -> None:
+    def setup(self, trainer, model: LightningModule) -> None:
         """
         Connects the plugins to the training process, creates optimizers
 
@@ -307,7 +306,7 @@ def on_train_end(self) -> None:
         """Hook to do something at the end of the training"""
         pass
 
-    def setup_optimizers(self, trainer: "Trainer"):
+    def setup_optimizers(self, trainer):
         """creates optimizers and schedulers
 
         Args:
diff --git a/pytorch_lightning/plugins/training_type/rpc_sequential.py b/pytorch_lightning/plugins/training_type/rpc_sequential.py
index b9a352aedb387..3878aa9db3ea4 100644
--- a/pytorch_lightning/plugins/training_type/rpc_sequential.py
+++ b/pytorch_lightning/plugins/training_type/rpc_sequential.py
@@ -21,7 +21,6 @@
 from torch.nn.parallel import DistributedDataParallel
 from torch.optim import Optimizer
 
-from pytorch_lightning.trainer.trainer import Trainer
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.overrides.distributed import LightningDistributedModule
 from pytorch_lightning.plugins.training_type.rpc import DEFAULT_RPC_TIMEOUT_SEC, RPCPlugin
@@ -330,11 +329,11 @@ def post_training_step(self):
         if self.main_rpc_process:
             super().post_training_step()
 
-    def start_training(self, trainer: 'Trainer') -> None:
+    def start_training(self, trainer) -> None:
         if self.main_rpc_process:
             super().start_training(trainer)
 
-    def start_testing(self, trainer: 'Trainer') -> None:
+    def start_testing(self, trainer) -> None:
         if self.main_rpc_process:
             super().start_testing(trainer)
 

From d9053174163f5835a7d6b20aa35e6f8cd10804cf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Sun, 14 Feb 2021 03:38:11 +0100
Subject: [PATCH 06/12] isort

---
 pytorch_lightning/accelerators/accelerator.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index 2e876d21d2481..d33c7506e165f 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -18,11 +18,7 @@
 from torch.utils.data import DataLoader
 
 from pytorch_lightning.core import LightningModule
-from pytorch_lightning.plugins.precision import (
-    ApexMixedPrecisionPlugin,
-    NativeMixedPrecisionPlugin,
-    PrecisionPlugin,
-)
+from pytorch_lightning.plugins.precision import ApexMixedPrecisionPlugin, NativeMixedPrecisionPlugin, PrecisionPlugin
 from pytorch_lightning.plugins.training_type import TrainingTypePlugin
 from pytorch_lightning.utilities.apply_func import move_data_to_device
 from pytorch_lightning.utilities.distributed import all_gather_ddp_if_available

From 5209f4704da616ef6fe97dd06dc7c88342adc6e4 Mon Sep 17 00:00:00 2001
From: Jirka Borovec <jirka.borovec@seznam.cz>
Date: Tue, 16 Feb 2021 23:49:57 +0100
Subject: [PATCH 07/12] fix

---
 pytorch_lightning/accelerators/accelerator.py | 17 -----------------
 1 file changed, 17 deletions(-)

diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index d33c7506e165f..d568fd525b25f 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -218,23 +218,6 @@ def validation_step_end(self, output):
         """
         return self.training_type_plugin.validation_step_end(output)
 
-    def predict(self, args):
-        """The prediction step.
-
-        Args:
-            args: the arguments for the models predict step. Can consist of the following:
-                batch (:class:`~torch.Tensor` | (:class:`~torch.Tensor`, ...) | [:class:`~torch.Tensor`, ...]):
-                    The output of your :class:`~torch.utils.data.DataLoader`. A tensor, tuple or list.
-                batch_idx (int): Integer displaying index of this batch
-                optimizer_idx (int): When using multiple optimizers, this argument will also be present.
-                hiddens(:class:`~torch.Tensor`): Passed in if
-                    :paramref:`~pytorch_lightning.trainer.trainer.Trainer.truncated_bptt_steps` > 0.
-
-        """
-        batch = self.to_device(args[0])
-        args[0] = batch
-        return self.training_type_plugin.predict(*args)
-
     def backward(
         self,
         closure_loss: torch.Tensor,

From 51861d3598ff85865149927edc848af987069624 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Thu, 18 Feb 2021 02:32:29 +0100
Subject: [PATCH 08/12] isort

---
 pytorch_lightning/plugins/training_type/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pytorch_lightning/plugins/training_type/__init__.py b/pytorch_lightning/plugins/training_type/__init__.py
index 06e3228219395..30723d67da3f4 100644
--- a/pytorch_lightning/plugins/training_type/__init__.py
+++ b/pytorch_lightning/plugins/training_type/__init__.py
@@ -1,6 +1,7 @@
 from pytorch_lightning.plugins.training_type.ddp import DDPPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.ddp2 import DDP2Plugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.ddp_spawn import DDPSpawnPlugin  # noqa: F401
+from pytorch_lightning.plugins.training_type.deepspeed import DeepSpeedPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.dp import DataParallelPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.horovod import HorovodPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin  # noqa: F401
@@ -12,4 +13,3 @@
 from pytorch_lightning.plugins.training_type.single_tpu import SingleTPUPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.tpu_spawn import TPUSpawnPlugin  # noqa: F401
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin  # noqa: F401
-from pytorch_lightning.plugins.training_type.deepspeed import DeepSpeedPlugin  # noqa: F401

From b8b1c7a9d9617b46c3b4d659a9c091cfc3e19f83 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Thu, 18 Feb 2021 02:35:10 +0100
Subject: [PATCH 09/12] resolve new flake8 errors

---
 pytorch_lightning/plugins/training_type/deepspeed.py | 2 +-
 pytorch_lightning/plugins/training_type/parallel.py  | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/pytorch_lightning/plugins/training_type/deepspeed.py b/pytorch_lightning/plugins/training_type/deepspeed.py
index 69fdb4c19a4b6..4bcceb518f1ea 100644
--- a/pytorch_lightning/plugins/training_type/deepspeed.py
+++ b/pytorch_lightning/plugins/training_type/deepspeed.py
@@ -263,7 +263,7 @@ def distributed_sampler_kwargs(self):
         distributed_sampler_kwargs = dict(num_replicas=self.world_size, rank=self.global_rank)
         return distributed_sampler_kwargs
 
-    def init_optimizers(self, trainer: "Trainer", model: LightningModule) -> Tuple[List, List, List]:
+    def init_optimizers(self, trainer, model: LightningModule) -> Tuple[List, List, List]:
         # Skip initializing optimizers here as DeepSpeed handles optimizers via config.
         # User may have specified config options instead in configure_optimizers, but this is handled
         # via `_initialize_deepspeed_train`
diff --git a/pytorch_lightning/plugins/training_type/parallel.py b/pytorch_lightning/plugins/training_type/parallel.py
index c27c9705cd362..a3f555e97f62d 100644
--- a/pytorch_lightning/plugins/training_type/parallel.py
+++ b/pytorch_lightning/plugins/training_type/parallel.py
@@ -24,7 +24,6 @@
 from pytorch_lightning.overrides.base import unwrap_lightning_module
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
-from pytorch_lightning.utilities import rank_zero_info
 from pytorch_lightning.utilities.distributed import all_gather_ddp_if_available, ReduceOp
 
 

From 6899965d4aa5d61b7f363d987f09865351f52d6b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Thu, 18 Feb 2021 02:43:15 +0100
Subject: [PATCH 10/12] fix unimported references of trainer

---
 pytorch_lightning/accelerators/accelerator.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index 2c25dffc5e91e..b7f029393d1d8 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -70,13 +70,13 @@ def setup(self, trainer, model: LightningModule) -> None:
         self.setup_optimizers(trainer)
         self.connect_precision_plugin(self.precision_plugin)
 
-    def start_training(self, trainer: 'Trainer'):
+    def start_training(self, trainer):
         self.training_type_plugin.start_training(trainer)
 
-    def start_testing(self, trainer: 'Trainer'):
+    def start_testing(self, trainer):
         self.training_type_plugin.start_testing(trainer)
 
-    def start_predicting(self, trainer: 'Trainer'):
+    def start_predicting(self, trainer):
         self.training_type_plugin.start_predicting(trainer)
 
     def pre_dispatch(self) -> None:

From 9d609fabdabf568c073d63bb6c81aedce3a258cf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Thu, 18 Feb 2021 16:07:08 +0100
Subject: [PATCH 11/12] pepe8

---
 pytorch_lightning/plugins/training_type/rpc.py       | 2 +-
 pytorch_lightning/plugins/training_type/tpu_spawn.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pytorch_lightning/plugins/training_type/rpc.py b/pytorch_lightning/plugins/training_type/rpc.py
index 5d41dbd7b4dcd..3c016f3cb8e25 100644
--- a/pytorch_lightning/plugins/training_type/rpc.py
+++ b/pytorch_lightning/plugins/training_type/rpc.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 import os
 from contextlib import suppress
-from typing import List, Optional, Sequence
+from typing import List, Optional
 
 import torch
 
diff --git a/pytorch_lightning/plugins/training_type/tpu_spawn.py b/pytorch_lightning/plugins/training_type/tpu_spawn.py
index 84d6de3827ebe..3fd2e0e659a9d 100644
--- a/pytorch_lightning/plugins/training_type/tpu_spawn.py
+++ b/pytorch_lightning/plugins/training_type/tpu_spawn.py
@@ -1,7 +1,7 @@
 import io
 import os
 import re
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Union
 
 import torch
 import torch.multiprocessing as mp

From df33f6927afb612659c28ca22fa1814c4d2fff04 Mon Sep 17 00:00:00 2001
From: tchaton <thomas@grid.ai>
Date: Thu, 18 Feb 2021 18:05:09 +0000
Subject: [PATCH 12/12] resolve tests

---
 tests/accelerators/test_tpu_backend.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/accelerators/test_tpu_backend.py b/tests/accelerators/test_tpu_backend.py
index daea22968b253..03da7c81b2b17 100644
--- a/tests/accelerators/test_tpu_backend.py
+++ b/tests/accelerators/test_tpu_backend.py
@@ -79,7 +79,7 @@ def test_if_test_works_after_train(tmpdir):
     model = BoringModel()
     trainer = Trainer(max_epochs=1, tpu_cores=8, default_root_dir=tmpdir, fast_dev_run=True)
     trainer.fit(model)
-    assert trainer.test(model) == 1
+    assert len(trainer.test(model)) == 1
 
 
 @pytest.mark.skipif(not _TPU_AVAILABLE, reason="test requires TPU machine")
@@ -119,4 +119,4 @@ def on_post_move_to_device(self):
         assert result
 
     assert not list(filter(lambda x: 'The model layers do not match' in str(x), warnings.list))
-    assert trainer.test(model) == 1
+    assert len(trainer.test(model)) == 1