code: def train_func(config):
      epochs = config.get("epochs", 3)
  
      model = resnet18()
  
      # Note that `prepare_model` needs to be called before setting optimizer.
      if not train.get_checkpoint():  # fresh start
          model = train.torch.prepare_model(model)
  
      # Create optimizer.
      optimizer_config = {
          "lr": config.get("lr"),
          "momentum": config.get("momentum"),
      }
      optimizer = torch.optim.SGD(
          model.parameters(), **optimizer_config
      )
  
      starting_epoch = 0
      if train.get_checkpoint():
          with train.get_checkpoint().as_directory() as checkpoint_dir:
              with open(
                  os.path.join(checkpoint_dir, "data.ckpt"), "rb"
              ) as fp:
                  checkpoint_dict = cpickle.load(fp)
  
          # Load in model
          model_state = checkpoint_dict["model"]
          model.load_state_dict(model_state)
          model = train.torch.prepare_model(model)
  
          # Load in optimizer
          optimizer_state = checkpoint_dict["optimizer_state_dict"]
          optimizer.load_state_dict(optimizer_state)
  
          # Optimizer configs (`lr`, `momentum`) are being mutated by PBT and passed in
          # through config, so we need to update the optimizer loaded from the checkpoint
          update_optimizer_config(optimizer, optimizer_config)
  
          # The current epoch increments the loaded epoch by 1
          checkpoint_epoch = checkpoint_dict["epoch"]
          starting_epoch = checkpoint_epoch + 1
  
      # Load in training and validation data.
      transform_train = transforms.Compose(
          [
              transforms.RandomCrop(32, padding=4),
              transforms.RandomHorizontalFlip(),
              transforms.ToTensor(),
              transforms.Normalize(
                  (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
              ),
          ]
      )  # meanstd transformation
  
      transform_test = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize(
                  (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
              ),
          ]
      )
  
      data_dir = config.get("data_dir", os.path.expanduser("~/data"))
      os.makedirs(data_dir, exist_ok=True)
      with FileLock(os.path.join(data_dir, ".ray.lock")):
          train_dataset = CIFAR10(
              root=data_dir,
              train=True,
              download=True,
              transform=transform_train,
          )
          validation_dataset = CIFAR10(
              root=data_dir,
              train=False,
              download=False,
              transform=transform_test,
          )
  
      if config.get("test_mode"):
          train_dataset = Subset(train_dataset, list(range(64)))
          validation_dataset = Subset(
              validation_dataset, list(range(64))
          )
  
      worker_batch_size = (
          config["batch_size"] // train.get_context().get_world_size()
      )
  
      train_loader = DataLoader(
          train_dataset, batch_size=worker_batch_size, shuffle=True
      )
      validation_loader = DataLoader(
          validation_dataset, batch_size=worker_batch_size
      )
  
      train_loader = train.torch.prepare_data_loader(train_loader)
      validation_loader = train.torch.prepare_data_loader(
          validation_loader
      )
  
      # Create loss.
      criterion = nn.CrossEntropyLoss()
  
      for epoch in range(starting_epoch, epochs):
          train_epoch(epoch, train_loader, model, criterion, optimizer)
          result = validate_epoch(validation_loader, model, criterion)
  
          with tempfile.TemporaryDirectory() as checkpoint_dir:
              with open(
                  os.path.join(checkpoint_dir, "data.ckpt"), "wb"
              ) as fp:
                  cpickle.dump(
                      {
                          "epoch": epoch,
                          "model": model.state_dict(),
                          "optimizer_state_dict": optimizer.state_dict(),
                      },
                      fp,
                  )
              checkpoint = Checkpoint.from_directory(checkpoint_dir)
              train.report(result, checkpoint=checkpoint)
  
imports: import argparse
  import os
  import tempfile
  import torch
  import torch.nn
  import torchvision.transforms
  from filelock import FileLock
  from torch.utils.data import DataLoader, Subset
  from torchvision.datasets import CIFAR10
  from torchvision.models import resnet18
  import ray
  import ray.cloudpickle
  from ray import train, tune
  from ray.train import Checkpoint, FailureConfig, RunConfig, ScalingConfig
  from ray.train.torch import TorchTrainer
  from ray.tune.schedulers import PopulationBasedTraining
  from ray.tune.tune_config import TuneConfig
  from ray.tune.tuner import Tuner
label_desc: Use image node_id 4 for calling a global function with example usage: train_func(config) without return types