code: def train_func(config):
      batch_size = config.get("batch_size", 32)
      hidden_size = config.get("hidden_size", 10)
      lr = config.get("lr", 1e-2)
      epochs = config.get("epochs", 3)
  
      train_dataset_shard = train.get_dataset_shard("train")
      validation_dataset = train.get_dataset_shard("validation")
  
      model = nn.Sequential(
          nn.Linear(100, hidden_size),
          nn.ReLU(),
          nn.Linear(hidden_size, 1),
      )
      model = train.torch.prepare_model(model)
  
      loss_fn = nn.L1Loss()
  
      optimizer = torch.optim.SGD(model.parameters(), lr=lr)
  
      results = []
  
      def create_torch_iterator(shard):
          iterator = shard.iter_torch_batches(batch_size=batch_size)
          for batch in iterator:
              yield batch["x"].float(), batch["y"].float()
  
      for _ in range(epochs):
          train_torch_dataset = create_torch_iterator(
              train_dataset_shard
          )
          validation_torch_dataset = create_torch_iterator(
              validation_dataset
          )
  
          device = train.torch.get_device()
  
          train_epoch(
              train_torch_dataset, model, loss_fn, optimizer, device
          )
          if train.get_context().get_world_rank() == 0:
              result = validate_epoch(
                  validation_torch_dataset, model, loss_fn, device
              )
          else:
              result = {}
          results.append(result)
  
          with tempfile.TemporaryDirectory() as tmpdir:
              torch.save(
                  model.module.state_dict(),
                  os.path.join(tmpdir, "model.pt"),
              )
              train.report(
                  result, checkpoint=Checkpoint.from_directory(tmpdir)
              )
  
      return results
  
file_path: ray/python/ray/train/examples/pytorch/torch_regression_example.py
imports: import argparse
  import os
  import tempfile
  from typing import Tuple
  import pandas
  import torch
  import torch.nn
  import ray
  import ray.train
  from ray.data import Dataset
  from ray.train import Checkpoint, DataConfig, ScalingConfig
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 4 for calling a global function with example usage: train_func(config) and returns: results