diff --git a/docs/source/usage_guides/memory.mdx b/docs/source/usage_guides/memory.mdx
index 8841cd584..213a2f670 100644
--- a/docs/source/usage_guides/memory.mdx
+++ b/docs/source/usage_guides/memory.mdx
@@ -25,16 +25,20 @@ training script. To use it, restructure your training function to include an inn
 and build your dataloaders inside it. At a minimum, this could look like 4 new lines of code. 
 > Note: The inner function *must* take in the batch size as the first parameter, but we do not pass one to it when called. The wrapper handles this for us
 
+It should also be noted that anything which will consume CUDA memory and passed to the `accelerator` **must** be declared inside the inner function,
+such as models and optimizers.
+
 ```diff
 def training_function(args):
     accelerator = Accelerator()
-    model = get_model()
-    model.to(accelerator.device)
-    optimizer = get_optimizer()
 
 +   @find_executable_batch_size(starting_batch_size=args.batch_size)
 +   def inner_training_loop(batch_size):
-+       nonlocal model, optimizer # Ensure they can be used in our context
++       nonlocal accelerator # Ensure they can be used in our context
++       accelerator.free_memory() # Free all lingering references
+        model = get_model()
+        model.to(accelerator.device)
+        optimizer = get_optimizer()
         train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
         lr_scheduler = get_scheduler(
             optimizer, 
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index 82dc34b8a..684a32853 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -122,24 +122,6 @@ def training_function(config, args):
 
     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)
-    # 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)
-
     # 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.
@@ -147,16 +129,31 @@ def training_function(config, args):
     @find_executable_batch_size(starting_batch_size=batch_size)
     def inner_training_loop(batch_size):
         # And now just move everything below under this function
-        # Ensure that anything declared outside this function is set as `nonlocal`
-        # so it is in scope
-        nonlocal model, optimizer
+        # 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) // gradient_accumulation_steps,
+            num_training_steps=(len(train_dataloader) * num_epochs),
         )
 
         # Prepare everything
@@ -174,12 +171,10 @@ def inner_training_loop(batch_size):
                 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()
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
 
             model.eval()
             for step, batch in enumerate(eval_dataloader):



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._