diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2d2a0bd56..1b3fa27e9 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -68,7 +68,7 @@
         DummyScheduler,
     )
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.distributed.xla_multiprocessing as xmp
 
 logger = get_logger(__name__)
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index 8b41c48f8..d5e816aab 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -33,7 +33,7 @@
 )
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 from .logging import get_logger
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index eb26c3a60..806fb7ddc 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -34,7 +34,7 @@
 )
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.distributed.parallel_loader as xpl
 
 
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
index 3d5977968..046d5ce5b 100644
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -21,7 +21,7 @@
 from .utils import DistributedType, honor_type, is_torch_version, is_tpu_available
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 36b11be56..63527124c 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -20,7 +20,7 @@
 from .utils import DistributedType, is_ccl_available, is_deepspeed_available, is_tpu_available
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 
@@ -88,7 +88,7 @@ def __init__(
                     "Please make sure to properly initialize your accelerator via `accelerator = Accelerator()` "
                     "before using any functionality from the `accelerate` library."
                 )
-            elif is_tpu_available() and not cpu:
+            if is_tpu_available() and not cpu:
                 self.distributed_type = DistributedType.TPU
                 self.num_processes = xm.xrt_world_size()
                 self.process_index = xm.get_ordinal()
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 594f7aa1c..a660db858 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -38,12 +38,7 @@
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
-    try:
-        # Will raise a RuntimeError if no XLA configuration is found
-        _ = xm.xla_device()
-        _tpu_available = True
-    except RuntimeError:
-        _tpu_available = False
+    _tpu_available = True
 except ImportError:
     _tpu_available = False
 
@@ -56,8 +51,15 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
-def is_tpu_available():
-    "Checks if `torch_xla` is installed and if a TPU is in the environment"
+def is_tpu_available(check_device=True):
+    "Checks if `torch_xla` is installed and potentially if a TPU is in the environment"
+    if _tpu_available and check_device:
+        try:
+            # Will raise a RuntimeError if no XLA configuration is found
+            _ = xm.xla_device()
+            return True
+        except RuntimeError:
+            return False
     return _tpu_available
 
 
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 954bb5dfd..46753513e 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -29,7 +29,7 @@
 from .versions import is_torch_version
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
index 0d4f5b7e2..206c70589 100644
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -28,7 +28,7 @@
 if is_deepspeed_available():
     from deepspeed import DeepSpeedEngine
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 
diff --git a/src/accelerate/utils/random.py b/src/accelerate/utils/random.py
index 3958022a4..51b7c3719 100644
--- a/src/accelerate/utils/random.py
+++ b/src/accelerate/utils/random.py
@@ -23,7 +23,7 @@
 from .imports import is_tpu_available
 
 
-if is_tpu_available():
+if is_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm
 
 


Review comments: body: Mmm, removing this will break the other places we use `is_tpu_available`
diff_hunk: @@ -38,12 +38,7 @@
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
-    try:
-        # Will raise a RuntimeError if no XLA configuration is found
-        _ = xm.xla_device()
-        _tpu_available = True
-    except RuntimeError:
-        _tpu_available = False
body: The added reason was specifically inside `AcceleratorState`, at the proposed location. But having it as an argument instead works as well. Will refactor
diff_hunk: @@ -38,12 +38,7 @@
 try:
     import torch_xla.core.xla_model as xm  # noqa: F401
 
-    try:
-        # Will raise a RuntimeError if no XLA configuration is found
-        _ = xm.xla_device()
-        _tpu_available = True
-    except RuntimeError:
-        _tpu_available = False
body: This needs to be `True` by default, and `False` when we don't want to check for the device (before launching multiprocessing for instance).
diff_hunk: @@ -56,8 +51,15 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
-def is_tpu_available():
-    "Checks if `torch_xla` is installed and if a TPU is in the environment"
+def is_tpu_available(check_device=False):
body: It needs to be `False` by default, because otherwise it also does this on import checks that are scattered around the library.
diff_hunk: @@ -56,8 +51,15 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
-def is_tpu_available():
-    "Checks if `torch_xla` is installed and if a TPU is in the environment"
+def is_tpu_available(check_device=False):
body: Or made into a separate function if we don't want `False` behavior. (I know we're not fans of that, but this is one case where it should be `False`)
diff_hunk: @@ -56,8 +51,15 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
-def is_tpu_available():
-    "Checks if `torch_xla` is installed and if a TPU is in the environment"
+def is_tpu_available(check_device=False):

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._