code: def training_function(config, args):
      # For testing only
      if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
          config["num_epochs"] = 2
      # 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")
  
      # New Code #
      # We now can define an inner training loop function. It should take a batch size as the only parameter,
      # and build the dataloaders in there.
      # It also gets our decorator
      @find_executable_batch_size(starting_batch_size=batch_size)
      def inner_training_loop(batch_size):
          # And now just move everything below under this function
          # We need to bring in the Accelerator object from earlier
          nonlocal accelerator
          # And reset all of its attributes that could hold onto any memory:
          accelerator.free_memory()
  
          # Then we can declare the model, optimizer, and everything else:
          set_seed(seed)
  
          # 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)
          train_dataloader, eval_dataloader = get_dataloaders(
              accelerator, batch_size
          )
  
          # Instantiate scheduler
          lr_scheduler = get_linear_schedule_with_warmup(
              optimizer=optimizer,
              num_warmup_steps=100,
              num_training_steps=(len(train_dataloader) * num_epochs),
          )
  
          # 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
                  accelerator.backward(loss)
                  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 #
      # And call it at the end with no arguments
      # Note: You could also refactor this outside of your training loop function
      inner_training_loop()
  
file_path: accelerate/examples/by_feature/memory.py
imports: import argparse
  import os
  import evaluate
  import torch
  from datasets import load_dataset
  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
  from accelerate.utils import find_executable_batch_size
label_desc: Use image node_id 2 for calling a global function with example usage: training_function(config, args) without return types