code: def train_loop_per_worker():
      hvd.init()
      dataset_shard = train.get_dataset_shard("train")
      model = NeuralNetwork()
      device = train.torch.get_device()
      model.to(device)
      loss_fn = nn.MSELoss()
      lr_scaler = 1
      optimizer = torch.optim.SGD(
          model.parameters(), lr=0.1 * lr_scaler
      )
      # Horovod: wrap optimizer with DistributedOptimizer.
      optimizer = hvd.DistributedOptimizer(
          optimizer,
          named_parameters=model.named_parameters(),
          op=hvd.Average,
      )
      for epoch in range(num_epochs):
          model.train()
          for batch in dataset_shard.iter_torch_batches(
              batch_size=32, dtypes=torch.float
          ):
              inputs, labels = (
                  torch.unsqueeze(batch["x"], 1),
                  batch["y"],
              )
              outputs = model(inputs)
              loss = loss_fn(outputs, labels)
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
              print(f"epoch: {epoch}, loss: {loss.item()}")
  
          with tempfile.TemporaryDirectory() as tmpdir:
              torch.save(
                  model.state_dict(), os.path.join(tmpdir, "model.pt")
              )
              train.report(
                  {"loss": loss.item()},
                  checkpoint=Checkpoint.from_directory(tmpdir),
              )
  
file_path: ray/doc/source/train/doc_code/hvd_trainer.py
imports: import os
  import tempfile
  import horovod.torch
  import ray
  from ray import train
  from ray.train import Checkpoint, ScalingConfig
  import ray.train.torch
  from ray.train.horovod import HorovodTrainer
  import torch
  import torch.nn
label_desc: Use image node_id 1 for calling a global function with example usage: train_loop_per_worker() without return types