code: def train_loop_per_worker(config):
      raw_model = resnet18(pretrained=True)
      model = train.torch.prepare_model(raw_model)
      criterion = nn.CrossEntropyLoss()
      optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
  
      train_dataset_shard = train.get_dataset_shard("train")
  
      for epoch in range(config["num_epochs"]):
          running_loss = 0.0
          for i, data in enumerate(
              train_dataset_shard.iter_torch_batches(
                  batch_size=config["batch_size"]
              )
          ):
              # get the inputs; data is a list of [inputs, labels]
              inputs = data["image"].to(device=train.torch.get_device())
              labels = data["label"].to(device=train.torch.get_device())
              # zero the parameter gradients
              optimizer.zero_grad()
  
              # forward + backward + optimize
              outputs = model(inputs)
              loss = criterion(outputs, labels)
              loss.backward()
              optimizer.step()
  
              # print statistics
              running_loss += loss.item()
              if i % 2000 == 1999:  # print every 2000 mini-batches
                  print(
                      f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}"
                  )
                  running_loss = 0.0
  
          with tempfile.TemporaryDirectory() as tmpdir:
              torch.save(
                  model.state_dict(), os.path.join(tmpdir, "model.pt")
              )
              train.report(
                  dict(running_loss=running_loss),
                  checkpoint=Checkpoint.from_directory(tmpdir),
              )
  
file_path: ray/release/air_tests/air_benchmarks/workloads/pytorch_training_e2e.py
imports: import click
  import time
  import json
  import os
  import tempfile
  from typing import Dict
  import numpy
  from torchvision import transforms
  from torchvision.models import resnet18
  import torch
  import torch.nn
  import torch.optim
  import ray
  from ray import train
  from ray.train import Checkpoint, RunConfig, ScalingConfig
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 3 for calling a global function with example usage: train_loop_per_worker(config) without return types