code: def train_loop_per_worker():
      worker_rank = train.get_context().get_world_rank()
      device = train.torch.get_device()
      world_size = ray.train.get_context().get_world_size()
      local_world_size = ray.train.get_context().get_local_world_size()
      torch_num_workers = args.torch_num_workers or os.cpu_count()
      # Divide by the number of Train workers because each has its own dataloader.
      torch_num_workers //= local_world_size
  
      # Setup the model
      raw_model = resnet50(weights=None)
      model = train.torch.prepare_model(raw_model)
      criterion = nn.CrossEntropyLoss()
      optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
  
      # Get the configured data loading solution.
      batch_iter = None
  
      if args.use_torch:
          batch_iter = get_torch_data_loader(
              worker_rank=worker_rank,
              batch_size=args.batch_size,
              num_workers=torch_num_workers,
              transform=get_transform(False),
          )
      elif args.use_mosaic:
          target_epoch_size = get_mosaic_epoch_size(
              args.num_workers, target_worker_gb=args.target_worker_gb
          )
          print(
              "Epoch size:",
              target_epoch_size
              if target_epoch_size is not None
              else "all",
              "images",
          )
          num_physical_nodes = world_size // local_world_size
          batch_iter = get_mosaic_dataloader(
              args.data_root,
              batch_size=args.batch_size,
              num_physical_nodes=num_physical_nodes,
              epoch_size=target_epoch_size,
              num_workers=torch_num_workers,
          )
  
      all_workers_time_list_across_epochs = []
      validation_accuracy_per_epoch = []
      # Validation loop with non-random cropped dataset
      # is only supported for image dataset.
      run_validation_set = (
          args.use_ray_data and args.file_type == "image"
      )
  
      # Begin training over the configured number of epochs.
      for epoch in range(args.num_epochs):
          # Ray Data needs to call iter_torch_batches on each epoch.
          if args.use_ray_data:
              batch_iter = _get_ray_data_batch_iterator(
                  args, worker_rank
              )
              if run_validation_set:
                  val_ds = train.get_dataset_shard("val")
                  batch_iter_val = val_ds.iter_torch_batches(
                      batch_size=args.batch_size
                  )
          # For synthetic data, we need to create the iterator each epoch.
          elif args.use_synthetic_data:
              # Generate a random batch, and continuously yield the same batch 1000 times.
              NUM_BATCHES_PER_EPOCH = 1000
              sample_batch = {
                  "image": torch.rand(
                      (
                          args.batch_size,
                          3,
                          DEFAULT_IMAGE_SIZE,
                          DEFAULT_IMAGE_SIZE,
                      ),
                      device=device,
                  ),
                  "label": torch.randint(
                      0,
                      NUM_BATCHES_PER_EPOCH,
                      (args.batch_size,),
                      device=device,
                  ),
              }
              batch_iter = itertools.repeat(
                  sample_batch, NUM_BATCHES_PER_EPOCH
              )
  
          print(f"Epoch {epoch+1} of {args.num_epochs}")
          num_rows = 0
          start_t = time.time()
          num_batches = 0.0
          total_loss = 0.0
          for batch_idx, batch in enumerate(batch_iter):
              num_rows += _get_batch_num_rows(batch)
  
              if not args.skip_train_model:
                  # get the inputs; data is a list of [inputs, labels]
                  inputs = torch.as_tensor(
                      batch["image"], dtype=torch.float32
                  ).to(device=device)
                  labels = torch.as_tensor(
                      batch["label"], dtype=torch.int64
                  ).to(device=device)
                  # forward + backward + optimize
                  outputs = model(inputs)
                  loss = criterion(outputs, labels)
  
                  # zero the parameter gradients
                  optimizer.zero_grad()
                  loss.backward()
                  optimizer.step()
  
                  num_batches += 1
                  total_loss += loss.item()
  
              # print statistics
              if (
                  batch_idx % 2000 == 1999
              ):  # print every 2000 mini-batches
                  print(
                      f"[{epoch + 1}, {batch_idx + 1:5d}]"
                      f"loss: {total_loss / 2000:.3f}"
                  )
          end_t = time.time()
  
          epoch_accuracy_val = None
          if run_validation_set and not args.skip_train_model:
              print(f"Starting validation set for epoch {epoch+1}")
              num_correct_val = 0
              num_rows_val = 0
              with torch.no_grad():
                  for batch_idx, batch in enumerate(batch_iter_val):
                      inputs = torch.as_tensor(
                          batch["image"], dtype=torch.float32
                      ).to(device=device)
                      labels = torch.as_tensor(
                          batch["label"], dtype=torch.int64
                      ).to(device=device)
  
                      outputs = model(inputs)
                      loss = criterion(outputs, labels)
                      output_classes = outputs.argmax(dim=1)
  
                      num_rows_val += len(labels)
                      num_correct_val += (
                          (output_classes == labels).sum().item()
                      )
              epoch_accuracy_val = num_correct_val / num_rows_val
              validation_accuracy_per_epoch.append(epoch_accuracy_val)
  
          with tempfile.TemporaryDirectory() as tmpdir:
              torch.save(
                  model.state_dict(), os.path.join(tmpdir, "model.pt")
              )
              checkpoint = Checkpoint.from_directory(tmpdir)
              train.report(
                  dict(
                      epoch_accuracy=epoch_accuracy_val,
                      loss_avg=(total_loss / num_batches)
                      if num_batches > 0
                      else 0,
                  ),
                  checkpoint=checkpoint,
              )
  
          # Workaround to report the epoch start/end time from each worker, so that we
          # can aggregate them at the end when calculating throughput.
          all_workers_time_list = [
              torch.zeros((2), dtype=torch.double, device=device)
              for _ in range(world_size)
          ]
          curr_worker_time = torch.tensor(
              [start_t, end_t], dtype=torch.double, device=device
          )
          dist.all_gather(all_workers_time_list, curr_worker_time)
          all_workers_time_list_across_epochs.append(
              all_workers_time_list
          )
  
          print(
              f"Epoch {epoch+1} of {args.num_epochs}, "
              f"tput: {num_rows / (end_t - start_t)}, "
              f"run time: {end_t - start_t}, "
              f"validation accuracy: "
              f"{epoch_accuracy_val * 100 if epoch_accuracy_val else 0:.3f}%"
          )
      # Similar reporting for aggregating number of rows across workers
      all_num_rows = [
          torch.zeros((1), dtype=torch.int32, device=device)
          for _ in range(world_size)
      ]
      curr_num_rows = torch.tensor(
          [num_rows], dtype=torch.int32, device=device
      )
      dist.all_gather(all_num_rows, curr_num_rows)
  
      per_epoch_times = {
          f"epoch_{i}_times": [
              tensor.tolist()
              for tensor in all_workers_time_list_across_epochs[i]
          ]
          for i in range(args.num_epochs)
      }
  
      final_train_report_metrics = {
          **per_epoch_times,
          "num_rows": [tensor.item() for tensor in all_num_rows],
      }
  
      if run_validation_set:
          all_num_rows_val = [
              torch.zeros((1), dtype=torch.int32, device=device)
              for _ in range(world_size)
          ]
          curr_num_rows_val = torch.tensor(
              [num_rows_val], dtype=torch.int32, device=device
          )
          dist.all_gather(all_num_rows_val, curr_num_rows_val)
  
          all_num_rows_correct_val = [
              torch.zeros((1), dtype=torch.int32, device=device)
              for _ in range(world_size)
          ]
          curr_num_rows_correct = torch.tensor(
              [num_correct_val], dtype=torch.int32, device=device
          )
          dist.all_gather(
              all_num_rows_correct_val, curr_num_rows_correct
          )
          final_train_report_metrics.update(
              {
                  "num_rows_val": [
                      tensor.item() for tensor in all_num_rows_val
                  ],
                  "num_rows_correct_val": [
                      tensor.item()
                      for tensor in all_num_rows_correct_val
                  ],
                  # Report the validation accuracy of the final epoch
                  "epoch_accuracy": validation_accuracy_per_epoch[-1],
              }
          )
  
      train.report(final_train_report_metrics)
  
imports: import ray
  from ray import train
  from ray.train import DataConfig, ScalingConfig, RunConfig, Checkpoint
  from ray.train.torch import TorchTrainer
  from ray.data.datasource.partitioning import Partitioning
  import tempfile
  import itertools
  import os
  import time
  from benchmark import Benchmark, BenchmarkMetric
  from image_loader_microbenchmark import get_transform, crop_and_flip_image, decode_image_crop_and_flip, center_crop_image
  from image_loader_microbenchmark import get_mosaic_dataloader
  import torch
  import torch.distributed
  from torchvision.models import resnet50
  from torchvision.transforms.functional import pil_to_tensor
  import torch.nn
  import torch.optim
  from dataset_benchmark_util import get_prop_parquet_paths, IMG_S3_ROOT, get_mosaic_epoch_size, IMAGENET_WNID_TO_ID
label_desc: Use image node_id 4 for calling a global function with example usage: train_loop_per_worker() without return types