code: def test_tune_torch_get_device_gpu(num_gpus_per_worker):
      """Tests if GPU ids are set correctly when running train concurrently in nested actors
      (for example when used with Tune).
      """
      from ray.train import ScalingConfig
  
      num_samples = 2
      num_workers = 2
  
      # We should have exactly enough resources in the cluster to run both samples
      # concurrently.
      total_gpus_required = (
          num_workers * num_gpus_per_worker * num_samples
      )
      # Divide by two because of a 2 node cluster.
      gpus_per_node = total_gpus_required // 2
  
      exception = None
      # Use the same number of cpus per node as gpus per node.
      with ray_start_2_node_cluster(
          num_cpus_per_node=gpus_per_node,
          num_gpus_per_node=gpus_per_node,
      ):
  
          @patch("torch.cuda.is_available", lambda: True)
          def train_fn():
              # We use STRICT_SPREAD strategy to force multiple samples on the same node.
              # For single or fractional GPU case, each worker has only 1 visible device (
              # the other is taken by the other sample) so device index should be 0.
              # For the multiple GPU case, each worker has 2 visible devices so device
              # index should be either 0 or 1. It doesn't matter which.
              devices = train.torch.get_device()
              if isinstance(devices, list):
                  assert sorted(
                      [device.index for device in devices]
                  ) == [0, 1]
              else:
                  assert train.torch.get_device().index == 0
  
          @ray.remote(num_cpus=0)
          class TrialActor:
              def __init__(self, warmup_steps):
                  self.trainer = TorchTrainer(
                      train_fn,
                      torch_config=TorchConfig(backend="gloo"),
                      run_config=RunConfig(
                          # Use a unique name to avoid using the same
                          # experiment directory
                          name=f"test_tune_torch_get_device_gpu_{uuid.uuid4()}"
                      ),
                      scaling_config=ScalingConfig(
                          num_workers=num_workers,
                          use_gpu=True,
                          resources_per_worker={
                              "CPU": 1,
                              "GPU": num_gpus_per_worker,
                          },
                          # Need to specify 0 trainer resources so STRICT_SPREAD
                          # will work.
                          trainer_resources={"CPU": 0},
                          placement_strategy="STRICT_SPREAD",
                          # Each gpu worker will be spread onto separate nodes. This
                          # forces different samples to run concurrently on the same
                          # node.
                      ),
                  )
  
              def run(self):
                  return self.trainer.fit()
  
          try:
              actors = [
                  TrialActor.remote(1) for _ in range(num_samples)
              ]
              ray.get([actor.run.remote() for actor in actors])
          except Exception as exc:
              exception = exc
  
      # Raise exception after Ray cluster has been shutdown to avoid corrupted state
      if exception:
          raise exception
  
imports: import contextlib
  import os
  import tempfile
  import time
  import uuid
  from unittest.mock import patch
  import pytest
  import torch
  import ray
  import ray.train
  from ray.cluster_utils import Cluster
  from ray.train import Checkpoint, RunConfig, ScalingConfig
  from ray.train.examples.pytorch.torch_linear_example import train_func
  from ray.train.torch import TorchCheckpoint, TorchConfig, TorchPredictor, TorchTrainer
  from ray.train.trainer import TrainingFailedError
label_desc: Use image node_id 10 for calling a global function with example usage: test_tune_torch_get_device_gpu(num_gpus_per_worker) and returns: self