code: def test_torch_conversion_collate_fn(ray_start_regular_shared):
      def collate_fn(batch: Dict[str, np.ndarray]):
          return torch.as_tensor(batch["id"] + 5)
  
      ds = ray.data.range(5)
      it = ds.iterator()
      for batch in it.iter_torch_batches(collate_fn=collate_fn):
          assert isinstance(batch, torch.Tensor)
          assert batch.tolist() == list(range(5, 10))
  
      # Should fail.
      with pytest.raises(ValueError):
          for batch in it.iter_torch_batches(
              collate_fn=collate_fn, dtypes=torch.float32
          ):
              assert isinstance(batch, torch.Tensor)
              assert batch.tolist() == list(range(5, 10))
  
      with pytest.raises(ValueError):
          for batch in it.iter_torch_batches(
              collate_fn=collate_fn, device="cpu"
          ):
              assert isinstance(batch, torch.Tensor)
              assert batch.tolist() == list(range(5, 10))
  
      # Test that we don't automatically set device if collate_fn is specified.
      with patch(
          "ray.air._internal.torch_utils.get_device",
          lambda: torch.device("cuda"),
      ):
          assert (
              ray.air._internal.torch_utils.get_device().type == "cuda"
          )
  
          it.iter_batches = MagicMock()
          for batch in it.iter_torch_batches(collate_fn=collate_fn):
              assert batch.device.type == "cpu"
              assert isinstance(batch, torch.Tensor)
              assert batch.tolist() == list(range(5, 10))
  
          # When collate_fn is specified, check that`_finalize_fn`
          # is not used in `DataIterator.iter_batches()`.
          iter_batches_calls_kwargs = [
              a.kwargs for a in it.iter_batches.call_args_list
          ]
          assert all(
              kwargs["_finalize_fn"] is None
              for kwargs in iter_batches_calls_kwargs
          ), iter_batches_calls_kwargs
  
file_path: ray/python/ray/data/tests/test_iterator.py
imports: from typing import Dict
  from unittest.mock import MagicMock, patch
  import numpy
  import pytest
  import tensorflow
  import torch
  import ray
label_desc: Use image node_id 10 for calling a global function with example usage: test_torch_conversion_collate_fn(ray_start_regular_shared) and returns: torch