code: def training_function(config, args):
      # New Code #
      test_predictions = []
      # Download the dataset
      datasets = load_dataset("glue", "mrpc")
      # Create our splits
      kfold = StratifiedKFold(n_splits=int(args.num_folds))
      # Initialize accelerator
      accelerator = Accelerator(
          cpu=args.cpu, mixed_precision=args.mixed_precision
      )
      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      lr = config["lr"]
      num_epochs = int(config["num_epochs"])
      seed = int(config["seed"])
      batch_size = int(config["batch_size"])
  
      metric = evaluate.load("glue", "mrpc")
  
      # If the batch size is too big we use gradient accumulation
      gradient_accumulation_steps = 1
      if (
          batch_size > MAX_GPU_BATCH_SIZE
          and accelerator.distributed_type != DistributedType.TPU
      ):
          gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
          batch_size = MAX_GPU_BATCH_SIZE
  
      set_seed(seed)
  
      # New Code #
      # Create our folds:
      folds = kfold.split(
          np.zeros(datasets["train"].num_rows),
          datasets["train"]["label"],
      )
      test_references = []
      # Iterate over them
      for i, (train_idxs, valid_idxs) in enumerate(folds):
          (
              train_dataloader,
              eval_dataloader,
              test_dataloader,
          ) = get_fold_dataloaders(
              accelerator,
              datasets,
              train_idxs,
              valid_idxs,
          )
          # Instantiate the model (we build the model here so that the seed also control new weights initialization)
          model = AutoModelForSequenceClassification.from_pretrained(
              "bert-base-cased", return_dict=True
          )
  
          # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
          # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
          # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
          model = model.to(accelerator.device)
  
          # Instantiate optimizer
          optimizer = AdamW(params=model.parameters(), lr=lr)
  
          # Instantiate scheduler
          lr_scheduler = get_linear_schedule_with_warmup(
              optimizer=optimizer,
              num_warmup_steps=100,
              num_training_steps=(len(train_dataloader) * num_epochs)
              // gradient_accumulation_steps,
          )
  
          # Prepare everything
          # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
          # prepare method.
          (
              model,
              optimizer,
              train_dataloader,
              eval_dataloader,
              lr_scheduler,
          ) = accelerator.prepare(
              model,
              optimizer,
              train_dataloader,
              eval_dataloader,
              lr_scheduler,
          )
  
          # Now we train the model
          for epoch in range(num_epochs):
              model.train()
              for step, batch in enumerate(train_dataloader):
                  # We could avoid this line since we set the accelerator with `device_placement=True`.
                  batch.to(accelerator.device)
                  outputs = model(**batch)
                  loss = outputs.loss
                  loss = loss / gradient_accumulation_steps
                  accelerator.backward(loss)
                  if step % gradient_accumulation_steps == 0:
                      optimizer.step()
                      lr_scheduler.step()
                      optimizer.zero_grad()
  
              model.eval()
              for step, batch in enumerate(eval_dataloader):
                  # We could avoid this line since we set the accelerator with `device_placement=True`.
                  batch.to(accelerator.device)
                  with torch.no_grad():
                      outputs = model(**batch)
                  predictions = outputs.logits.argmax(dim=-1)
                  (
                      predictions,
                      references,
                  ) = accelerator.gather_for_metrics(
                      (predictions, batch["labels"])
                  )
                  metric.add_batch(
                      predictions=predictions,
                      references=references,
                  )
  
              eval_metric = metric.compute()
              # Use accelerator.print to print only on the main process.
              accelerator.print(f"epoch {epoch}:", eval_metric)
  
          # New Code #
          # We also run predictions on the test set at the very end
          fold_predictions = []
          for step, batch in enumerate(test_dataloader):
              # We could avoid this line since we set the accelerator with `device_placement=True`.
              batch.to(accelerator.device)
              with torch.no_grad():
                  outputs = model(**batch)
              predictions = outputs.logits
              predictions, references = accelerator.gather_for_metrics(
                  (predictions, batch["labels"])
              )
              fold_predictions.append(predictions.cpu())
              if i == 0:
                  # We need all of the test predictions
                  test_references.append(references.cpu())
          # Use accelerator.print to print only on the main process.
          test_predictions.append(torch.cat(fold_predictions, dim=0))
          # We now need to release all our memory and get rid of the current model, optimizer, etc
          accelerator.free_memory()
      # New Code #
      # Finally we check the accuracy of our folded results:
      test_references = torch.cat(test_references, dim=0)
      preds = (
          torch.stack(test_predictions, dim=0)
          .sum(dim=0)
          .div(int(args.num_folds))
          .argmax(dim=-1)
      )
      test_metric = metric.compute(
          predictions=preds, references=test_references
      )
      accelerator.print(
          "Average test metrics from all folds:", test_metric
      )
  
file_path: accelerate/examples/by_feature/cross_validation.py
imports: import argparse
  from typing import List
  import evaluate
  import numpy
  import torch
  from datasets import DatasetDict, load_dataset
  from sklearn.model_selection import StratifiedKFold
  from torch.optim import AdamW
  from torch.utils.data import DataLoader
  from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
  from accelerate import Accelerator, DistributedType
label_desc: Use image node_id 2 for calling a global function with example usage: training_function(config, args) without return types