code: def test_torch_linear_failure(ray_start_4_cpus):
      num_workers = 2
      epochs = 3
  
      from ray.train.torch import TorchConfig
  
      test_config = TorchConfig()
  
      train_func = functools.partial(
          linear_train_func,
          {"lr": 1e-3, "batch_size": 64, "epochs": epochs},
      )
  
      iterator = create_iterator(
          train_func, test_config, num_workers=num_workers
      )
      kill_callback = KillCallback(
          fail_on=1, backend_executor=iterator._backend_executor
      )
  
      for intermediate_result in iterator:
          assert len(intermediate_result) == num_workers
          kill_callback.handle_result(intermediate_result)
  
      results = kill_callback.results
      assert len(results) == epochs
      for i in range(num_workers):
          last_result = results[-1][i].metrics
          first_result = results[0][i].metrics
          assert last_result["loss"] < first_result["loss"]
  
imports: import functools
  import time
  from unittest.mock import patch
  import pytest
  import ray
  from ray import train
  from ray.air._internal.util import StartTraceback
  from ray.train import DataConfig
  from ray.train._internal.backend_executor import BackendExecutor
  from ray.train._internal.session import get_session, init_session
  from ray.train._internal.utils import construct_train_func
  from ray.train._internal.worker_group import WorkerGroup
  from ray.train.backend import BackendConfig
  from ray.train.examples.pytorch.torch_linear_example import train_func
  from ray.train.examples.tf.tensorflow_mnist_example import train_func
  from ray.train.tests.util import mock_storage_context
  from ray.train.trainer import TrainingIterator
label_desc: Use image node_id 16 for calling a global function with example usage: test_torch_linear_failure(ray_start_4_cpus) without return types