code: def test_torch_iter_torch_batches_auto_device(
      ray_start_4_cpus_2_gpus, use_gpu
  ):
      """
      Tests that iter_torch_batches in TorchTrainer worker function uses the
      default device.
      """
  
      def train_fn():
          dataset = train.get_dataset_shard("train")
          for batch in dataset.iter_torch_batches(
              dtypes=torch.float, device="cpu"
          ):
              assert str(batch["data"].device) == "cpu"
  
          # Autodetect
          for batch in dataset.iter_torch_batches(dtypes=torch.float):
              assert str(batch["data"].device) == str(
                  train.torch.get_device()
              )
  
      dataset = ray.data.from_numpy(
          np.array([[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]]).T
      )
      # Test that this works outside a Train function
      for batch in dataset.iter_torch_batches(
          dtypes=torch.float, device="cpu"
      ):
          assert str(batch["data"].device) == "cpu"
  
      trainer = TorchTrainer(
          train_fn,
          scaling_config=ScalingConfig(num_workers=2, use_gpu=use_gpu),
          datasets={"train": dataset},
      )
      trainer.fit()
  
file_path: ray/python/ray/train/tests/test_gpu_2.py
imports: import numpy
  import pytest
  import torch
  import ray
  import ray.data
  import ray.train
  from ray import tune
  from ray.air.config import ScalingConfig
  from ray.train.examples.pytorch.torch_linear_example import LinearDataset
  from ray.train.torch.torch_trainer import TorchTrainer
label_desc: Use image node_id 1 for calling a global function with example usage: test_torch_iter_torch_batches_auto_device(ray_start_4_cpus_2_gpus, use_gpu) without return types