code: def train_func(config):
      data_size = config.get("data_size", 1000)
      val_size = config.get("val_size", 400)
      batch_size = config.get("batch_size", 32)
      hidden_size = config.get("hidden_size", 1)
      lr = config.get("lr", 1e-2)
      epochs = config.get("epochs", 3)
  
      train_dataset = LinearDataset(2, 5, size=data_size)
      val_dataset = LinearDataset(2, 5, size=val_size)
      train_loader = torch.utils.data.DataLoader(
          train_dataset, batch_size=batch_size
      )
      validation_loader = torch.utils.data.DataLoader(
          val_dataset, batch_size=batch_size
      )
  
      train_loader = train.torch.prepare_data_loader(train_loader)
      validation_loader = train.torch.prepare_data_loader(
          validation_loader
      )
  
      model = nn.Linear(1, hidden_size)
      model = train.torch.prepare_model(model)
  
      loss_fn = nn.MSELoss()
  
      optimizer = torch.optim.SGD(model.parameters(), lr=lr)
  
      results = []
      for epoch in range(epochs):
          train_epoch(epoch, train_loader, model, loss_fn, optimizer)
          state_dict, loss = validate_epoch(
              validation_loader, model, loss_fn
          )
          result = dict(loss=loss)
          results.append(result)
  
          with tempfile.TemporaryDirectory() as tmpdir:
              torch.save(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_linear_example.py
imports: import argparse
  import os
  import tempfile
  import numpy
  import torch
  import torch.nn
  import ray.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_func(config) and returns: results