code: def train_loop_per_worker(config):
      import horovod.torch as hvd
  
      hvd.init()
      device = ray.train.torch.get_device()
      net = resnet18().to(device)
      optimizer = torch.optim.SGD(
          net.parameters(),
          lr=config["lr"],
      )
      epoch = 0
  
      checkpoint = train.get_checkpoint()
      if checkpoint:
          with checkpoint.as_directory() as checkpoint_dir:
              checkpoint_dir = Path(checkpoint_dir)
              model_state = torch.load(
                  checkpoint_dir / "model.pt", map_location="cpu"
              )
              optimizer_state = torch.load(
                  checkpoint_dir / "optim.pt", map_location="cpu"
              )
              epoch = (
                  torch.load(checkpoint_dir / "extra_state.pt")["epoch"]
                  + 1
              )
  
          net.load_state_dict(model_state)
          optimizer.load_state_dict(optimizer_state)
  
      criterion = nn.CrossEntropyLoss()
      optimizer = hvd.DistributedOptimizer(optimizer)
      np.random.seed(1 + hvd.rank())
      torch.manual_seed(1234)
      # To ensure consistent initialization across workers,
      hvd.broadcast_parameters(net.state_dict(), root_rank=0)
      hvd.broadcast_optimizer_state(optimizer, root_rank=0)
  
      trainset = ray.get(config["data"])
  
      train_sampler = torch.utils.data.distributed.DistributedSampler(
          trainset, num_replicas=hvd.size(), rank=hvd.rank()
      )
  
      # Note, don't set `num_workers` in DataLoader (not even 1),
      # as that will separately start multiple processes (each corresponding to 1 worker)
      # to load the data. This is known to cause issues with Ray.
      trainloader = DataLoader(
          trainset,
          batch_size=int(config["batch_size"]),
          sampler=train_sampler,
      )
  
      for epoch in range(
          epoch, 40
      ):  # loop over the dataset multiple times
          running_loss = 0.0
          epoch_steps = 0
          for i, data in enumerate(trainloader):
              # get the inputs; data is a list of [inputs, labels]
              inputs, labels = data
              inputs, labels = inputs.to(device), labels.to(device)
  
              # zero the parameter gradients
              optimizer.zero_grad()
  
              # forward + backward + optimize
              outputs = net(inputs)
              loss = criterion(outputs, labels)
              loss.backward()
              optimizer.step()
  
              # print statistics
              running_loss += loss.item()
              epoch_steps += 1
  
              if i % 2000 == 1999:  # print every 2000 mini-batches
                  print(
                      "[%d, %5d] loss: %.3f"
                      % (epoch + 1, i + 1, running_loss / epoch_steps)
                  )
  
              if config["smoke_test"]:
                  break
  
          with tempfile.TemporaryDirectory() as tmpdir:
              torch.save(
                  net.state_dict(), os.path.join(tmpdir, "model.pt")
              )
              torch.save(
                  optimizer.state_dict(),
                  os.path.join(tmpdir, "optim.pt"),
              )
              torch.save(
                  {"epoch": epoch},
                  os.path.join(tmpdir, "extra_state.pt"),
              )
              train.report(
                  dict(loss=running_loss / epoch_steps),
                  checkpoint=Checkpoint.from_directory(tmpdir),
              )
  
file_path: ray/release/air_tests/horovod/workloads/horovod_tune_test.py
imports: import os
  from pathlib import Path
  import tempfile
  import numpy
  import torch
  import torch.nn
  from torch.utils.data import DataLoader
  import torchvision
  import torchvision.transforms
  from torchvision.models import resnet18
  import ray
  from ray.train import Checkpoint, CheckpointConfig, FailureConfig, RunConfig, ScalingConfig
  import ray.train.torch
  from ray.train.horovod import HorovodTrainer
  from ray import train, tune
  from ray.tune.schedulers import create_scheduler
  from ray.tune.tune_config import TuneConfig
  from ray.tune.tuner import Tuner
  from ray.tune.utils.release_test_util import ProgressCallback
label_desc: Use image node_id 1 for calling a global function with example usage: train_loop_per_worker(config) without return types