code: def linear_train_func(accelerator: Accelerator, config):
      from accelerate.utils import DummyOptim
      from deepspeed.ops.adam import DeepSpeedCPUAdam
  
      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
      )
  
      model = nn.Linear(1, hidden_size)
  
      loss_fn = nn.MSELoss()
      if (
          accelerator.state.deepspeed_plugin
          and "optimizer"
          in accelerator.state.deepspeed_plugin.deepspeed_config
      ):
          optimizer_cls = DummyOptim
      elif accelerator.state.deepspeed_plugin:
          optimizer_cls = DeepSpeedCPUAdam
      else:
          optimizer_cls = torch.optim.SGD
  
      # Accelerate boilerplate
      no_decay = ["bias", "LayerNorm.weight"]
      optimizer_grouped_parameters = [
          {
              "params": [
                  p
                  for n, p in model.named_parameters()
                  if not any(nd in n for nd in no_decay)
              ],
              "weight_decay": 0.0,
          },
          {
              "params": [
                  p
                  for n, p in model.named_parameters()
                  if any(nd in n for nd in no_decay)
              ],
              "weight_decay": 0.0,
          },
      ]
      optimizer = optimizer_cls(optimizer_grouped_parameters, lr=lr)
      (
          train_loader,
          validation_loader,
          model,
          optimizer,
      ) = accelerator.prepare(
          train_loader, validation_loader, model, optimizer
      )
  
      results = []
      for _ in range(epochs):
          for X, y in train_loader:
              # Compute prediction error
              pred = model(X)
              loss = loss_fn(pred, y)
  
              # Backpropagation
              accelerator.backward(loss)
              optimizer.step()
              optimizer.zero_grad()
  
          num_batches = len(validation_loader)
          model.eval()
          loss = 0
          with torch.no_grad():
              for X, y in validation_loader:
                  pred = model(X)
                  loss += loss_fn(pred, y).item()
          loss /= num_batches
          import copy
  
          model_copy = copy.deepcopy(accelerator.unwrap_model(model))
          state_dict, loss = model_copy.cpu().state_dict(), loss
  
          result = dict(loss=loss)
          results.append(result)
  
          with TemporaryDirectory() as tmpdir:
              torch.save(
                  state_dict, os.path.join(tmpdir, "checkpoint.pt")
              )
              train.report(
                  result, checkpoint=Checkpoint.from_directory(tmpdir)
              )
  
      return results
  
file_path: ray/python/ray/train/tests/test_torch_accelerate.py
imports: import os
  from tempfile import TemporaryDirectory
  import pytest
  import torch
  import torch.nn
  from accelerate import Accelerator
  import ray
  import ray.train
  from ray.train import Checkpoint, ScalingConfig
  from ray.train.examples.pytorch.torch_linear_example import LinearDataset
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 2 for calling a global function with example usage: linear_train_func(accelerator, config) and returns: results