aval

Lightning-AI · Dec 2, 2020 · 9c0c567 · 9c0c567
1 parent b34d615
commit 9c0c567
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 30 deletions.
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
@@ -20,7 +20,6 @@
 import os
 import re
 import tempfile
-import types
 from abc import ABC
 from argparse import Namespace
 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
@@ -36,7 +35,7 @@
 from pytorch_lightning.core.memory import ModelSummary
 from pytorch_lightning.core.saving import ALLOWED_CONFIG_TYPES, PRIMITIVE_TYPES, ModelIO
 from pytorch_lightning.core.step_result import Result
-from pytorch_lightning.utilities import TPU_AVAILABLE, AMPType, rank_zero_warn, XLADeviceUtils
+from pytorch_lightning.utilities import TPU_AVAILABLE, rank_zero_warn
 from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.parsing import AttributeDict, collect_init_args, get_init_args

diff --git a/pytorch_lightning/core/optimizer.py b/pytorch_lightning/core/optimizer.py
@@ -11,38 +11,14 @@
 # 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 collections
-import copy
-import inspect
-import os
-import re
-import tempfile
 import types
-from abc import ABC
-from argparse import Namespace
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Optional
 from weakref import proxy
 
-import torch
-from torch import ScriptModule, Tensor
-from torch.nn import Module
-from torch.optim import SGD
 from torch.optim.optimizer import Optimizer
 
-from pytorch_lightning import _logger as log
-from pytorch_lightning.callbacks import Callback
-from pytorch_lightning.core.grads import GradInformation
-from pytorch_lightning.core.hooks import CheckpointHooks, DataHooks, ModelHooks
-from pytorch_lightning.core.memory import ModelSummary
-from pytorch_lightning.core.saving import ALLOWED_CONFIG_TYPES, PRIMITIVE_TYPES, ModelIO
-from pytorch_lightning.core.step_result import Result
-from pytorch_lightning.utilities import AMPType, rank_zero_warn
-from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
+from pytorch_lightning.utilities import TPU_AVAILABLE
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.parsing import AttributeDict, collect_init_args, get_init_args
-from pytorch_lightning.utilities.xla_device_utils import XLADeviceUtils
-
-TPU_AVAILABLE = XLADeviceUtils.tpu_device_exists()
 
 if TPU_AVAILABLE:
     import torch_xla.core.xla_model as xm

diff --git a/pytorch_lightning/utilities/device_parser.py b/pytorch_lightning/utilities/device_parser.py
@@ -13,8 +13,9 @@
 # limitations under the License.
 import torch
 from typing import Union, Any, List, Optional, MutableSequence
+
+from pytorch_lightning.utilities import TPU_AVAILABLE
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.xla_device_utils import XLADeviceUtils
 
 
 def determine_root_gpu_device(gpus: List[int]) -> Optional[int]:
@@ -105,7 +106,7 @@ def parse_tpu_cores(tpu_cores: Union[int, str, List]) -> Optional[Union[List[int
     if not _tpu_cores_valid(tpu_cores):
         raise MisconfigurationException("`tpu_cores` can only be 1, 8 or [<1-8>]")
 
-    if tpu_cores is not None and not XLADeviceUtils.tpu_device_exists():
+    if tpu_cores is not None and not TPU_AVAILABLE:
         raise MisconfigurationException('No TPU devices were found.')
 
     return tpu_cores