diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 1b21ebed7..5e364491e 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -107,6 +107,9 @@ def collate_fn(examples):
 
 
 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
diff --git a/examples/by_feature/fsdp_with_peak_mem_tracking.py b/examples/by_feature/fsdp_with_peak_mem_tracking.py
index 09199154a..8046b4557 100644
--- a/examples/by_feature/fsdp_with_peak_mem_tracking.py
+++ b/examples/by_feature/fsdp_with_peak_mem_tracking.py
@@ -83,6 +83,9 @@ def __exit__(self, *exc):
 
 
 def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
     # Initialize accelerator
     if args.with_tracking:
         accelerator = Accelerator(
diff --git a/examples/by_feature/gradient_accumulation.py b/examples/by_feature/gradient_accumulation.py
index d95f9bef3..fe3d948f5 100644
--- a/examples/by_feature/gradient_accumulation.py
+++ b/examples/by_feature/gradient_accumulation.py
@@ -102,6 +102,9 @@ def collate_fn(examples):
 
 
 def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
     # New Code #
     gradient_accumulation_steps = int(args.gradient_accumulation_steps)
     # Initialize accelerator
diff --git a/examples/by_feature/memory.py b/examples/by_feature/memory.py
index b11c80fd0..f71283561 100644
--- a/examples/by_feature/memory.py
+++ b/examples/by_feature/memory.py
@@ -107,6 +107,9 @@ def collate_fn(examples):
 
 
 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
diff --git a/examples/by_feature/multi_process_metrics.py b/examples/by_feature/multi_process_metrics.py
index dc47f3db6..d69c0bc90 100644
--- a/examples/by_feature/multi_process_metrics.py
+++ b/examples/by_feature/multi_process_metrics.py
@@ -109,6 +109,9 @@ def collate_fn(examples):
 
 
 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
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index 00c8665dd..d453b770c 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -107,6 +107,9 @@ def collate_fn(examples):
 
 
 def training_function(config, args):
+    # For testing only
+    if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
+        config["num_epochs"] = 2
     # Initialize Accelerator
 
     # New Code #
diff --git a/src/accelerate/test_utils/examples.py b/src/accelerate/test_utils/examples.py
index 4e4092c0e..f459e03c4 100644
--- a/src/accelerate/test_utils/examples.py
+++ b/src/accelerate/test_utils/examples.py
@@ -115,11 +115,17 @@ def compare_against_test(base_filename: str, feature_filename: str, parser_only:
     # Specific code in our script that differs from the full version, aka what is new
     new_feature_code = []
     passed_idxs = []  # We keep track of the idxs just in case it's a repeated statement
-    for i, line in enumerate(feature_file_func):
+    it = iter(feature_file_func)
+    for i in range(len(feature_file_func) - 1):
         if i not in passed_idxs:
+            line = next(it)
             if (line not in full_file_func) and (line.lstrip() != _dl_line):
-                new_feature_code.append(line)
-                passed_idxs.append(i)
+                if "TESTING_MOCKED_DATALOADERS" not in line:
+                    new_feature_code.append(line)
+                    passed_idxs.append(i)
+                else:
+                    # Skip over the `config['num_epochs'] = 2` statement
+                    _ = next(it)
 
     # Extract out just the new parts from the full_file_training_func
     new_full_example_parts = []
@@ -127,8 +133,9 @@ def compare_against_test(base_filename: str, feature_filename: str, parser_only:
     for i, line in enumerate(base_file_func):
         if i not in passed_idxs:
             if (line not in full_file_func) and (line.lstrip() != _dl_line):
-                new_full_example_parts.append(line)
-                passed_idxs.append(i)
+                if "TESTING_MOCKED_DATALOADERS" not in line:
+                    new_full_example_parts.append(line)
+                    passed_idxs.append(i)
 
     # Finally, get the overall diff
     diff_from_example = [line for line in new_feature_code if line not in new_full_example_parts]
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 9e1854cd8..361e19642 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -20,8 +20,6 @@
 import unittest
 from unittest import mock
 
-import torch
-
 from accelerate.test_utils.examples import compare_against_test
 from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
 from accelerate.utils import write_basic_config
@@ -145,7 +143,7 @@ def test_checkpointing_by_epoch(self):
         --output_dir {self.tmpdir}
         """.split()
         run_command(self._launch_args + testargs)
-        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_1")))
+        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "epoch_0")))
 
     def test_checkpointing_by_steps(self):
         testargs = f"""
@@ -154,17 +152,16 @@ def test_checkpointing_by_steps(self):
         --output_dir {self.tmpdir}
         """.split()
         _ = run_command(self._launch_args + testargs)
-        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_5")))
+        self.assertTrue(os.path.exists(os.path.join(self.tmpdir, "step_2")))
 
     def test_load_states_by_epoch(self):
         testargs = f"""
         examples/by_feature/checkpointing.py
-        --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_1")}
+        --resume_from_checkpoint {os.path.join(self.tmpdir, "epoch_0")}
         """.split()
         output = run_command(self._launch_args + testargs, return_stdout=True)
         self.assertNotIn("epoch 0:", output)
-        self.assertNotIn("epoch 1:", output)
-        self.assertIn("epoch 2:", output)
+        self.assertIn("epoch 1:", output)
 
     def test_load_states_by_steps(self):
         testargs = f"""
@@ -172,17 +169,8 @@ def test_load_states_by_steps(self):
         --resume_from_checkpoint {os.path.join(self.tmpdir, "step_5")}
         """.split()
         output = run_command(self._launch_args + testargs, return_stdout=True)
-        if torch.cuda.is_available():
-            num_processes = torch.cuda.device_count()
-        else:
-            num_processes = 1
-        if num_processes > 1:
-            self.assertNotIn("epoch 0:", output)
-            self.assertNotIn("epoch 1:", output)
-        else:
-            self.assertNotIn("epoch 0:", output)
-            self.assertIn("epoch 1:", output)
-        self.assertIn("epoch 2:", output)
+        self.assertNotIn("epoch 0:", output)
+        self.assertIn("epoch 1:", output)
 
     @slow
     def test_cross_validation(self):



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