add warning message.

Author: vanbasten23

torch_xla/core/xla_model.py

@@ -5,6 +5,7 @@
 import re
 import threading
 import time
+import warnings
 from typing import List, Optional
 import torch
 import torch.distributed._functional_collectives
@@ -88,6 +89,11 @@ def get_xla_supported_devices(devkind=None, max_devices=None):
  Returns:
  The list of device strings.
  """
+ # TODO(xiowei): Remove the below if statement after r2.2 release.
+ if devkind.casefold() == 'gpu':
+ warnings.warn("GPU as a device name is being deprecate. Please replace it with CUDA such as get_xla_supported_devices(devkind='CUDA'). Similarly, please replace PJRT_DEVICE=GPU with PJRT_DEVICE=CUDA.")
+ devkind = 'CUDA'
+
  xla_devices = _DEVICES.value
  devkind = [devkind] if devkind else [
  'TPU', 'GPU', 'XPU', 'NEURON', 'CPU', 'CUDA', 'ROCM'

torch_xla/runtime.py

@@ -107,6 +107,11 @@ def xla_device(n: Optional[int] = None,
  Returns:
  A `torch.device` representing an XLA device.
  """
+ # TODO(xiowei): Remove the warning message at r2.2 release.
+ pjrt_device = xu.getenv_as(xenv.PJRT_DEVICE, str)
+ if pjrt_device.casefold() == 'gpu':
+ warnings.warn('PJRT_DEVICE=GPU is being deprecate. Please replace PJRT_DEVICE=GPU with PJRT_DEVICE=CUDA.')
+
  if n is None:
  return torch.device(torch_xla._XLAC._xla_get_default_device())