fix comment and test failure.

Author: vanbasten23

.circleci/common.sh

@@ -131,6 +131,8 @@ function run_torch_xla_python_tests() {
  if [ -x "$(command -v nvidia-smi)" ]; then
  # These tests fail on CUDA with 03/30 TF-pin update (https://github.com/pytorch/xla/pull/4840)
  PJRT_DEVICE=CUDA python test/test_train_mp_imagenet_fsdp.py --fake_data --use_nested_fsdp --use_small_fake_sample --num_epochs=1
+ # TODO(xiowei replace gpu with cuda): remove the test below with PJRT_DEVICE=GPU because PJRT_DEVICE=GPU is being deprecated.
+ PJRT_DEVICE=GPU python test/test_train_mp_imagenet_fsdp.py --fake_data --use_nested_fsdp --use_small_fake_sample --num_epochs=1
  PJRT_DEVICE=CUDA python test/test_train_mp_imagenet_fsdp.py --fake_data --auto_wrap_policy type_based --use_small_fake_sample --num_epochs=1
  XLA_DISABLE_FUNCTIONALIZATION=1 PJRT_DEVICE=CUDA python test/test_train_mp_imagenet_fsdp.py --fake_data --use_nested_fsdp --use_small_fake_sample --num_epochs=1
  # Syncfree SGD optimizer tests

torch_xla/core/xla_model.py

@@ -89,9 +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.")
+ # TODO(xiowei replace gpu with cuda): Remove the below if statement after r2.2 release.
+ if devkind and 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

torch_xla/runtime.py

@@ -107,10 +107,12 @@ 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.
+ # TODO(xiowei replace gpu with cuda): 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.')
+ 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())