update tests

Author: awaelchli

tests/accelerators/test_gpu.py

@@ -11,10 +11,7 @@
 def test_get_torch_gpu_stats(tmpdir):
     """Test GPU get_device_stats with Pytorch >= 1.8.0."""
     current_device = torch.device(f"cuda:{torch.cuda.current_device()}")
-    GPUAccel = GPUAccelerator(
-        training_type_plugin=DataParallelPlugin(parallel_devices=[current_device]), precision_plugin=PrecisionPlugin()
-    )
-    gpu_stats = GPUAccel.get_device_stats(current_device)
+    gpu_stats = GPUAccelerator().get_device_stats(current_device)
     fields = ["allocated_bytes.all.freed", "inactive_split.all.peak", "reserved_bytes.large_pool.peak"]
 
     for f in fields:
@@ -26,10 +23,7 @@ def test_get_torch_gpu_stats(tmpdir):
 def test_get_nvidia_gpu_stats(tmpdir):
     """Test GPU get_device_stats with Pytorch < 1.8.0."""
     current_device = torch.device(f"cuda:{torch.cuda.current_device()}")
-    GPUAccel = GPUAccelerator(
-        training_type_plugin=DataParallelPlugin(parallel_devices=[current_device]), precision_plugin=PrecisionPlugin()
-    )
-    gpu_stats = GPUAccel.get_device_stats(current_device)
+    gpu_stats = GPUAccelerator().get_device_stats(current_device)
     fields = ["utilization.gpu", "memory.used", "memory.free", "utilization.memory"]
 
     for f in fields:

tests/accelerators/test_tpu.py

@@ -13,7 +13,7 @@
 # limitations under the License
 import collections
 from copy import deepcopy
-from unittest.mock import patch
+from unittest.mock import patch, Mock
 
 import pytest
 import torch
@@ -288,13 +288,13 @@ def forward(self, x):
 
 
 def test_tpu_invalid_raises():
-    accelerator = TPUAccelerator(object(), TPUSpawnPlugin())
+    training_type_plugin = TPUSpawnPlugin(accelerator=TPUAccelerator(), precision_plugin=Mock())
     with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `TPUPrecisionPlugin"):
-        training_type_plugin.setup(object())
+        training_type_plugin.setup(Mock())
 
-    accelerator = TPUAccelerator(TPUPrecisionPlugin(), DDPPlugin())
+    training_type_plugin = DDPPlugin(accelerator=TPUAccelerator(), precision_plugin=TPUPrecisionPlugin())
     with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUPlugin` or `TPUSpawnPlugi"):
-        training_type_plugin.setup(object())
+        training_type_plugin.setup(Mock())
 
 
 def test_tpu_invalid_raises_set_precision_with_strategy():