diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index c213a9167..5020c6cc4 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -14,4 +14,17 @@ jobs:
     - name: Install Python dependencies
       run: pip install setuptools==59.5.0; pip install -e .[test,test_trackers]
     - name: Run Tests
-      run: make test
\ No newline at end of file
+      run: make test
+      
+  test_examples:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.6
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.6
+    - name: Install Python dependencies
+      run: pip install setuptools==59.5.0; pip install -e .[test] tensorboard
+    - name: Run Tests
+      run: make test_examples
\ No newline at end of file
diff --git a/Makefile b/Makefile
index e5b557f10..db7893565 100644
--- a/Makefile
+++ b/Makefile
@@ -25,4 +25,7 @@ style:
 	
 # Run tests for the library
 test:
-	python -m pytest -n auto --dist=loadfile -s -v ./tests/
\ No newline at end of file
+	python -m pytest -n auto --dist=loadfile -s -v ./tests/ --ignore=./tests/test_examples.py
+
+test_examples:
+	python -m pytest -n auto --dist=loadfile -s -v ./tests/test_examples.py
diff --git a/examples/by_feature/checkpointing.py b/examples/by_feature/checkpointing.py
index 35205107d..e595a8db7 100644
--- a/examples/by_feature/checkpointing.py
+++ b/examples/by_feature/checkpointing.py
@@ -115,14 +115,17 @@ def training_function(config, args):
 
     # New Code #
     # Parse out whether we are saving every epoch or after a certain number of batches
-    if args.checkpointing_steps == "epoch":
-        checkpointing_steps = args.checkpointing_steps
-    elif args.checkpointing_steps.isdigit():
-        checkpointing_steps = int(args.checkpointing_steps)
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
     else:
-        raise ValueError(
-            f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
-        )
+        checkpointing_steps = None
 
     set_seed(seed)
 
@@ -162,21 +165,29 @@ def training_function(config, args):
 
     # New Code #
     # We need to keep track of how many total steps we have iterated over
-    if isinstance(checkpointing_steps, int):
-        overall_step = 0
+    overall_step = 0
 
     # We need to load the checkpoint back in before training here with `load_state`
     # The total number of epochs is adjusted based on where the state is being loaded from,
     # as we assume continuation of the same training script
     if args.resume_from_checkpoint:
-        accelerator.print(f"Resuming from checkpoint: {args.resume_from_checkpoint}")
-        accelerator.load_state(args.resume_from_checkpoint)
-
-        if "epoch" in args.resume_from_checkpoint:
-            num_epochs -= int(args.resume_from_checkpoint.replace("epoch_", ""))
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            num_epochs -= int(training_difference.replace("epoch_", ""))
             resume_step = None
         else:
-            resume_step = int(args.resume_from_checkpoint.replace("step_", ""))
+            resume_step = int(training_difference.replace("step_", ""))
             num_epochs -= resume_step // len(train_dataloader)
             # If resuming by step, we also need to know exactly how far into the DataLoader we went
             resume_step = (num_epochs * len(train_dataloader)) - resume_step
@@ -200,6 +211,8 @@ def training_function(config, args):
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
+            # New Code #
+            overall_step += 1
 
             # New Code #
             # We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
@@ -237,7 +250,7 @@ def training_function(config, args):
         # Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
         # If mixed precision was used, will also save a "scalar.bin" file
         if checkpointing_steps == "epoch":
-            output_dir = f"epoch_{num_epochs}"
+            output_dir = f"epoch_{epoch}"
             if args.output_dir is not None:
                 output_dir = os.path.join(args.output_dir, output_dir)
             accelerator.save_state(output_dir)
@@ -258,7 +271,7 @@ def main():
     parser.add_argument(
         "--checkpointing_steps",
         type=str,
-        default="epoch",
+        default=None,
         help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
     )
     parser.add_argument(
diff --git a/examples/by_feature/tracking.py b/examples/by_feature/tracking.py
index ca11fa04b..170211e5e 100644
--- a/examples/by_feature/tracking.py
+++ b/examples/by_feature/tracking.py
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
+import logging
+import os
 
 import torch
 from torch.utils.data import DataLoader
@@ -28,6 +30,9 @@
 )
 
 
+logger = logging.getLogger(__name__)
+
+
 ########################################################################
 # This is a fully working simple example to use Accelerate,
 # specifically showcasing the experiment tracking capability,
@@ -110,7 +115,9 @@ def training_function(config, args):
     # Note: If using a custom `Tracker` class, should be passed in here such as:
     # >>> log_with = ["all", MyCustomTrackerClassInstance()]
     if args.with_tracking:
-        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", logging_dir=args.logging_dir
+        )
     else:
         accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
@@ -158,7 +165,10 @@ def training_function(config, args):
     # New Code #
     # We need to initalize the trackers we use. Overall configurations can also be stored
     if args.with_tracking:
-        accelerator.init_trackers("accelerate_glue_with_tracking", config)
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+        accelerator.init_trackers(run, config)
 
     # Now we train the model
     for epoch in range(num_epochs):
@@ -173,7 +183,8 @@ def training_function(config, args):
             outputs = model(**batch)
             loss = outputs.loss
             # New Code #
-            total_loss += loss.detach().float()
+            if args.with_tracking:
+                total_loss += loss.detach().float()
             loss = loss / gradient_accumulation_steps
             accelerator.backward(loss)
             if step % gradient_accumulation_steps == 0:
@@ -202,14 +213,21 @@ def training_function(config, args):
         # New Code #
         # To actually log, we call `Accelerator.log`
         # The values passed can be of `str`, `int`, or `float`
-        accelerator.log(
-            {"accuracy": eval_metric["accuracy"], "f1": eval_metric["f1"], "train_loss": total_loss, "epoch": epoch}
-        )
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "accuracy": eval_metric["accuracy"],
+                    "f1": eval_metric["f1"],
+                    "train_loss": total_loss,
+                    "epoch": epoch,
+                }
+            )
 
     # New Code #
     # When a run is finished, you should call `accelerator.end_training()`
     # to close all of the open trackers
-    accelerator.end_training()
+    if args.with_tracking:
+        accelerator.end_training()
 
 
 def main():
@@ -229,6 +247,12 @@ def main():
         action="store_true",
         help="Whether to load in all available experiment trackers from the environment and use them for logging.",
     )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
     training_function(config, args)
diff --git a/examples/complete_cv_example.py b/examples/complete_cv_example.py
index e62568a39..01033498c 100644
--- a/examples/complete_cv_example.py
+++ b/examples/complete_cv_example.py
@@ -74,16 +74,12 @@ def __getitem__(self, idx):
 def training_function(config, args):
     # Initialize accelerator
     if args.with_tracking:
-        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", logging_dir=args.logging_dir
+        )
     else:
-        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
 
-    if hasattr(args.checkpointing_steps, "isdigit"):
-        checkpointing_steps = args.checkpointing_steps
-        if args.checkpointing_steps.isdigit():
-            checkpointing_steps = int(args.checkpointing_steps)
-    else:
-        checkpointing_steps = None
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
     num_epochs = int(config["num_epochs"])
@@ -93,9 +89,25 @@ def training_function(config, args):
     if not isinstance(image_size, (list, tuple)):
         image_size = (image_size, image_size)
 
+    # Parse out whether we are saving every epoch or after a certain number of batches
+    if hasattr(args.checkpointing_steps, "isdigit"):
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
+            checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
+    else:
+        checkpointing_steps = None
+
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
-        accelerator.init_trackers("cv_example", config)
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+        accelerator.init_trackers(run, config)
 
     # Grab all the image filenames
     file_names = [os.path.join(args.data_dir, fname) for fname in os.listdir(args.data_dir) if fname.endswith(".jpg")]
@@ -163,24 +175,27 @@ def training_function(config, args):
     )
 
     # Potentially load in the weights and states from a previous save
-    state_restored = True
     if args.resume_from_checkpoint:
         if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
             accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
             accelerator.load_state(args.resume_from_checkpoint)
-            resume_step = None
+            path = os.path.basename(args.resume_from_checkpoint)
         else:
             # Get the most recent checkpoint
             dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
             dirs.sort(key=os.path.getctime)
             path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
-            if "epoch" in path.name:
-                num_epochs -= int(path.name.replace("epoch_", ""))
-            else:
-                resume_step = int(path.name.replace("step_", ""))
-                num_epochs -= resume_step // len(train_dataloader)
-                resume_step = (num_epochs * len(train_dataloader)) - resume_step
-                state_restored = False
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            num_epochs -= int(training_difference.replace("epoch_", ""))
+            resume_step = None
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            num_epochs -= resume_step // len(train_dataloader)
+            # If resuming by step, we also need to know exactly how far into the DataLoader we went
+            resume_step = (num_epochs * len(train_dataloader)) - resume_step
 
     overall_step = 0
     # Now we train the model
@@ -190,8 +205,9 @@ def training_function(config, args):
             total_loss = 0
         for step, batch in enumerate(train_dataloader):
             # We need to skip steps until we reach the resumed step
-            if args.resume_from_checkpoint and epoch == 0 and step < resume_step:
-                continue
+            if args.resume_from_checkpoint and epoch == 0:
+                if resume_step is not None and step < resume_step:
+                    pass
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch = {k: v.to(accelerator.device) for k, v in batch.items()}
             inputs = (batch["image"] - mean) / std
@@ -206,32 +222,40 @@ def training_function(config, args):
             optimizer.zero_grad()
             overall_step += 1
             if isinstance(checkpointing_steps, int):
+                output_dir = f"step_{overall_step}"
                 if overall_step % checkpointing_steps == 0:
-                    accelerator.save_state(f"step_{overall_step}")
-        if state_restored:
-            model.eval()
-            accurate = 0
-            num_elems = 0
-            for step, batch in enumerate(eval_dataloader):
-                # We could avoid this line since we set the accelerator with `device_placement=True`.
-                batch = {k: v.to(accelerator.device) for k, v in batch.items()}
-                inputs = (batch["image"] - mean) / std
-                with torch.no_grad():
-                    outputs = model(inputs)
-                predictions = outputs.argmax(dim=-1)
-                accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
-                num_elems += accurate_preds.shape[0]
-                accurate += accurate_preds.long().sum()
-
-            eval_metric = accurate.item() / num_elems
-            # Use accelerator.print to print only on the main process.
-            accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
-            if args.with_tracking:
-                accelerator.log(
-                    {"accuracy": 100 * eval_metric, "total_loss": total_loss, "epoch": epoch}, step=overall_step
-                )
-            if args.checkpointing_steps == "epoch":
-                accelerator.save_state(f"epoch_{epoch}")
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+        model.eval()
+        accurate = 0
+        num_elems = 0
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch = {k: v.to(accelerator.device) for k, v in batch.items()}
+            inputs = (batch["image"] - mean) / std
+            with torch.no_grad():
+                outputs = model(inputs)
+            predictions = outputs.argmax(dim=-1)
+            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
+            num_elems += accurate_preds.shape[0]
+            accurate += accurate_preds.long().sum()
+
+        eval_metric = accurate.item() / num_elems
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")
+        if args.with_tracking:
+            accelerator.log(
+                {"accuracy": 100 * eval_metric, "total_loss": total_loss, "epoch": epoch}, step=overall_step
+            )
+        if checkpointing_steps == "epoch":
+            output_dir = f"epoch_{epoch}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
+
+    if args.with_tracking:
+        accelerator.end_training()
 
 
 def main():
@@ -254,6 +278,12 @@ def main():
         default=None,
         help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
     )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
     parser.add_argument(
         "--resume_from_checkpoint",
         type=str,
@@ -262,9 +292,15 @@ def main():
     )
     parser.add_argument(
         "--with_tracking",
-        required=False,
+        action="store_true",
         help="Whether to load in all available experiment trackers from the environment and use them for logging.",
     )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
     args = parser.parse_args()
     config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}
     training_function(config, args)
diff --git a/examples/complete_nlp_example.py b/examples/complete_nlp_example.py
index 57125aa29..a0e8d568c 100644
--- a/examples/complete_nlp_example.py
+++ b/examples/complete_nlp_example.py
@@ -55,14 +55,21 @@
 def training_function(config, args):
     # Initialize accelerator
     if args.with_tracking:
-        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
+        accelerator = Accelerator(
+            cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all", logging_dir=args.logging_dir
+        )
     else:
-        accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+        accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
 
     if hasattr(args.checkpointing_steps, "isdigit"):
-        checkpointing_steps = args.checkpointing_steps
-        if args.checkpointing_steps.isdigit():
+        if args.checkpointing_steps == "epoch":
+            checkpointing_steps = args.checkpointing_steps
+        elif args.checkpointing_steps.isdigit():
             checkpointing_steps = int(args.checkpointing_steps)
+        else:
+            raise ValueError(
+                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
+            )
     else:
         checkpointing_steps = None
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
@@ -74,7 +81,10 @@ def training_function(config, args):
 
     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
-        accelerator.init_trackers("nlp_example", config)
+        run = os.path.split(__file__)[-1].split(".")[0]
+        if args.logging_dir:
+            run = os.path.join(args.logging_dir, run)
+        accelerator.init_trackers(run, config)
 
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     datasets = load_dataset("glue", "mrpc")
@@ -143,27 +153,31 @@ def collate_fn(examples):
         model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
 
+    overall_step = 0
+
     # Potentially load in the weights and states from a previous save
-    state_restored = True
     if args.resume_from_checkpoint:
         if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
             accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
             accelerator.load_state(args.resume_from_checkpoint)
-            resume_step = None
+            path = os.path.basename(args.resume_from_checkpoint)
         else:
             # Get the most recent checkpoint
             dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
             dirs.sort(key=os.path.getctime)
             path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
-            if "epoch" in path.name:
-                num_epochs -= int(path.name.replace("epoch_", ""))
-            else:
-                resume_step = int(path.name.replace("step_", ""))
-                num_epochs -= resume_step // len(train_dataloader)
-                resume_step = (num_epochs * len(train_dataloader)) - resume_step
-                state_restored = False
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            num_epochs -= int(training_difference.replace("epoch_", ""))
+            resume_step = None
+        else:
+            resume_step = int(training_difference.replace("step_", ""))
+            num_epochs -= resume_step // len(train_dataloader)
+            # If resuming by step, we also need to know exactly how far into the DataLoader we went
+            resume_step = (num_epochs * len(train_dataloader)) - resume_step
 
-    overall_step = 0
     # Now we train the model
     for epoch in range(num_epochs):
         model.train()
@@ -171,8 +185,9 @@ def collate_fn(examples):
             total_loss = 0
         for step, batch in enumerate(train_dataloader):
             # We need to skip steps until we reach the resumed step
-            if args.resume_from_checkpoint and epoch == 0 and step < resume_step:
-                continue
+            if args.resume_from_checkpoint and epoch == 0:
+                if resume_step is not None and step < resume_step:
+                    pass
             # We could avoid this line since we set the accelerator with `device_placement=True`.
             batch.to(accelerator.device)
             outputs = model(**batch)
@@ -190,42 +205,51 @@ def collate_fn(examples):
             overall_step += 1
 
             if isinstance(checkpointing_steps, int):
+                output_dir = f"step_{overall_step}"
                 if overall_step % checkpointing_steps == 0:
-                    accelerator.save_state(f"step_{overall_step}")
-        if state_restored:
-            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)
-                metric.add_batch(
-                    predictions=accelerator.gather(predictions),
-                    references=accelerator.gather(batch["labels"]),
-                )
-
-            eval_metric = metric.compute()
-            # Use accelerator.print to print only on the main process.
-            accelerator.print(f"epoch {epoch}:", eval_metric)
-            if args.with_tracking:
-                accelerator.log(
-                    {
-                        "accuracy": eval_metric["accuracy"],
-                        "f1": eval_metric["f1"],
-                        "total_loss": total_loss,
-                        "epoch": epoch,
-                    },
-                    step=overall_step,
-                )
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+
+        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)
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather((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)
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "accuracy": eval_metric["accuracy"],
+                    "f1": eval_metric["f1"],
+                    "train_loss": total_loss,
+                    "epoch": epoch,
+                }
+            )
+
+        if checkpointing_steps == "epoch":
+            output_dir = f"epoch_{epoch}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
 
-            if args.checkpointing_steps == "epoch":
-                accelerator.save_state(f"epoch_{epoch}")
+    if args.with_tracking:
+        accelerator.end_training()
 
 
 def main():
     parser = argparse.ArgumentParser(description="Simple example of training script.")
-    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
     parser.add_argument(
         "--mixed_precision",
         type=str,
@@ -250,9 +274,21 @@ def main():
     )
     parser.add_argument(
         "--with_tracking",
-        required=False,
+        action="store_true",
         help="Whether to load in all available experiment trackers from the environment and use them for logging.",
     )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help="Location on where to store experiment tracking logs`",
+    )
     args = parser.parse_args()
     config = {"lr": 2e-5, "num_epochs": 3, "correct_bias": True, "seed": 42, "batch_size": 16}
     training_function(config, args)
diff --git a/examples/cv_example.py b/examples/cv_example.py
index b14b067be..a80b42d13 100644
--- a/examples/cv_example.py
+++ b/examples/cv_example.py
@@ -73,7 +73,7 @@ def __getitem__(self, idx):
 
 def training_function(config, args):
     # Initialize accelerator
-    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mix_precision)
+    accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mix_precision)
 
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     lr = config["lr"]
@@ -185,7 +185,6 @@ def training_function(config, args):
 def main():
     parser = argparse.ArgumentParser(description="Simple example of training script.")
     parser.add_argument("--data_dir", required=True, help="The data folder on disk.")
-    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
     parser.add_argument(
         "--mixed_precision",
         type=str,
@@ -195,6 +194,12 @@ def main():
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",
     )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
     parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
     args = parser.parse_args()
     config = {"lr": 3e-2, "num_epochs": 3, "seed": 42, "batch_size": 64, "image_size": 224}
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
index 915d35807..87cfd7698 100644
--- a/examples/nlp_example.py
+++ b/examples/nlp_example.py
@@ -49,21 +49,19 @@
 EVAL_BATCH_SIZE = 32
 
 
-def training_function(config, args):
-    # 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"])
-    correct_bias = config["correct_bias"]
-    seed = int(config["seed"])
-    batch_size = int(config["batch_size"])
-
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset,
+    using "bert-base-cased" as the tokenizer.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+    """
     tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
     datasets = load_dataset("glue", "mrpc")
-    metric = load_metric("glue", "mrpc")
-
-    set_seed(seed)
 
     def tokenize_function(examples):
         # max_length=None => use the model max length (it's actually the default)
@@ -81,12 +79,6 @@ def tokenize_function(examples):
     # transformers library
     tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
 
-    # If the batch size is too big we use gradient accumulation
-    gradient_accumulation_steps = 1
-    if batch_size > MAX_GPU_BATCH_SIZE:
-        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
-        batch_size = MAX_GPU_BATCH_SIZE
-
     def collate_fn(examples):
         # On TPU it's best to pad everything to the same length or training will be very slow.
         if accelerator.distributed_type == DistributedType.TPU:
@@ -101,6 +93,29 @@ def collate_fn(examples):
         tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
     )
 
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # 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"])
+    correct_bias = config["correct_bias"]
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+
+    metric = load_metric("glue", "mrpc")
+
+    # If the batch size is too big we use gradient accumulation
+    gradient_accumulation_steps = 1
+    if batch_size > MAX_GPU_BATCH_SIZE:
+        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
+        batch_size = MAX_GPU_BATCH_SIZE
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)
     # 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)
 
@@ -148,9 +163,10 @@ def collate_fn(examples):
             with torch.no_grad():
                 outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather((predictions, batch["labels"]))
             metric.add_batch(
-                predictions=accelerator.gather(predictions),
-                references=accelerator.gather(batch["labels"]),
+                predictions=predictions,
+                references=references,
             )
 
         eval_metric = metric.compute()
diff --git a/setup.py b/setup.py
index 0ce1e5af4..7e26759d1 100644
--- a/setup.py
+++ b/setup.py
@@ -20,7 +20,12 @@
 extras["docs"] = []
 extras["test"] = [
     "pytest",
-    "pytest-xdist"
+    "pytest-xdist",
+    "pytest-subtests",
+    "datasets",
+    "transformers",
+    "scipy",
+    "sklearn"
 ]
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorflow"]
 extras["dev"] = extras["quality"] + extras["test"]
diff --git a/src/accelerate/test_utils/examples.py b/src/accelerate/test_utils/examples.py
new file mode 100644
index 000000000..4e4092c0e
--- /dev/null
+++ b/src/accelerate/test_utils/examples.py
@@ -0,0 +1,139 @@
+#!/usr/bin/env python
+
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+A collection of utilities for comparing `examples/complete_*_example.py` scripts with the capabilities inside of each
+`examples/by_feature` example. `compare_against_test` is the main function that should be used when testing, while the
+others are used to either get the code that matters, or to preprocess them (such as stripping comments)
+"""
+
+import os
+from typing import List
+
+
+def get_function_contents_by_name(lines: List[str], name: str):
+    """
+    Extracts a function from `lines` of segmented source code with the name `name`.
+
+    Args:
+        lines (`List[str]`):
+            Source code of a script seperated by line.
+        name (`str`):
+            The name of the function to extract. Should be either `training_function` or `main`
+    """
+    if name != "training_function" and name != "main":
+        raise ValueError(f"Incorrect function name passed: {name}, choose either 'main' or 'training_function'")
+    good_lines, found_start = [], False
+    for line in lines:
+        if not found_start and f"def {name}" in line:
+            found_start = True
+            good_lines.append(line)
+            continue
+        if found_start:
+            if name == "training_function" and "def main" in line:
+                return good_lines
+            if name == "main" and "if __name__" in line:
+                return good_lines
+            good_lines.append(line)
+
+
+def clean_lines(lines: List[str]):
+    """
+    Filters `lines` and removes any entries that start with a comment ('#') or is just a newline ('\n')
+
+    Args:
+        lines (`List[str]`):
+            Source code of a script seperated by line.
+    """
+    return [line for line in lines if not line.lstrip().startswith("#") and line != "\n"]
+
+
+def compare_against_test(base_filename: str, feature_filename: str, parser_only: bool, secondary_filename: str = None):
+    """
+    Tests whether the additional code inside of `feature_filename` was implemented in `base_filename`. This should be
+    used when testing to see if `complete_*_.py` examples have all of the implementations from each of the
+    `examples/by_feature/*` scripts.
+
+    It utilizes `nlp_example.py` to extract out all of the repeated training code, so that only the new additional code
+    is examined and checked. If something *other* than `nlp_example.py` should be used, such as `cv_example.py` for the
+    `complete_cv_example.py` script, it should be passed in for the `secondary_filename` parameter.
+
+    Args:
+        base_filename (`str` or `os.PathLike`):
+            The filepath of a single "complete" example script to test, such as `examples/complete_cv_example.py`
+        feature_filename (`str` or `os.PathLike`):
+            The filepath of a single feature example script. The contents of this script are checked to see if they
+            exist in `base_filename`
+        parser_only (`bool`):
+            Whether to compare only the `main()` sections in both files, or to compare the contents of
+            `training_loop()`
+        secondary_filename (`str`, *optional*):
+            A potential secondary filepath that should be included in the check. This function extracts the base
+            functionalities off of "examples/nlp_example.py", so if `base_filename` is a script other than
+            `complete_nlp_example.py`, the template script should be included here. Such as `examples/cv_example.py`
+    """
+    with open(base_filename, "r") as f:
+        base_file_contents = f.readlines()
+    with open(os.path.abspath(os.path.join("examples", "nlp_example.py")), "r") as f:
+        full_file_contents = f.readlines()
+    with open(feature_filename, "r") as f:
+        feature_file_contents = f.readlines()
+    if secondary_filename is not None:
+        with open(secondary_filename, "r") as f:
+            secondary_file_contents = f.readlines()
+
+    # This is our base, we remove all the code from here in our `full_filename` and `feature_filename` to find the new content
+    if parser_only:
+        base_file_func = clean_lines(get_function_contents_by_name(base_file_contents, "main"))
+        full_file_func = clean_lines(get_function_contents_by_name(full_file_contents, "main"))
+        feature_file_func = clean_lines(get_function_contents_by_name(feature_file_contents, "main"))
+        if secondary_filename is not None:
+            secondary_file_func = clean_lines(get_function_contents_by_name(secondary_file_contents, "main"))
+    else:
+        base_file_func = clean_lines(get_function_contents_by_name(base_file_contents, "training_function"))
+        full_file_func = clean_lines(get_function_contents_by_name(full_file_contents, "training_function"))
+        feature_file_func = clean_lines(get_function_contents_by_name(feature_file_contents, "training_function"))
+        if secondary_filename is not None:
+            secondary_file_func = clean_lines(
+                get_function_contents_by_name(secondary_file_contents, "training_function")
+            )
+
+    _dl_line = "train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n"
+
+    # 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):
+        if i not in passed_idxs:
+            if (line not in full_file_func) and (line.lstrip() != _dl_line):
+                new_feature_code.append(line)
+                passed_idxs.append(i)
+
+    # Extract out just the new parts from the full_file_training_func
+    new_full_example_parts = []
+    passed_idxs = []  # We keep track of the idxs just in case it's a repeated statement
+    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)
+
+    # Finally, get the overall diff
+    diff_from_example = [line for line in new_feature_code if line not in new_full_example_parts]
+    if secondary_filename is not None:
+        diff_from_two = [line for line in full_file_contents if line not in secondary_file_func]
+        diff_from_example = [line for line in diff_from_example if line not in diff_from_two]
+
+    return diff_from_example
diff --git a/tests/test_examples.py b/tests/test_examples.py
new file mode 100644
index 000000000..d60d7c08f
--- /dev/null
+++ b/tests/test_examples.py
@@ -0,0 +1,190 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+import tempfile
+import unittest
+from unittest import mock
+
+from torch.utils.data import DataLoader
+
+from accelerate import DistributedType
+from accelerate.test_utils.examples import compare_against_test
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+SRC_DIRS = [os.path.abspath(os.path.join("examples", "by_feature"))]
+sys.path.extend(SRC_DIRS)
+
+if SRC_DIRS is not None:
+    import checkpointing
+    import tracking
+
+# DataLoaders built from `test_samples/MRPC` for quick testing
+# Should mock `{script_name}.get_dataloaders` via:
+# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
+
+
+def mocked_dataloaders(accelerator, batch_size: int = 16):
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
+    datasets = load_dataset("csv", data_files=data_files)
+    label_list = datasets["train"].unique("label")
+
+    label_to_id = {v: i for i, v in enumerate(label_list)}
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(
+            examples["sentence1"], examples["sentence2"], truncation=True, max_length=None, padding="max_length"
+        )
+        if "label" in examples:
+            outputs["labels"] = [label_to_id[l] for l in examples["label"]]
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["sentence1", "sentence2", "label"],
+    )
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=2)
+    eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1)
+
+    return train_dataloader, eval_dataloader
+
+
+class ExampleDifferenceTests(unittest.TestCase):
+    """
+    This TestCase checks that all of the `complete_*` scripts contain all of the
+    information found in the `by_feature` scripts, line for line. If one fails,
+    then a complete example does not contain all of the features in the features
+    scripts, and should be updated.
+
+    Each example script should be a single test (such as `test_nlp_example`),
+    and should run `one_complete_example` twice: once with `parser_only=True`,
+    and the other with `parser_only=False`. This is so that when the test
+    failures are returned to the user, they understand if the discrepancy lies in
+    the `main` function, or the `training_loop` function. Otherwise it will be
+    unclear.
+
+    Also, if there are any expected differences between the base script used and
+    `complete_nlp_example.py` (the canonical base script), these should be included in
+    `special_strings`. These would be differences in how something is logged, print statements,
+    etc (such as calls to `Accelerate.log()`)
+    """
+
+    def one_complete_example(
+        self, complete_file_name: str, parser_only: bool, secondary_filename: str = None, special_strings: list = None
+    ):
+        """
+        Tests a single `complete` example against all of the implemented `by_feature` scripts
+
+        Args:
+            complete_file_name (`str`):
+                The filename of a complete example
+            parser_only (`bool`):
+                Whether to look at the main training function, or the argument parser
+            secondary_filename (`str`, *optional*):
+                A potential secondary base file to strip all script information not relevant for checking,
+                such as "cv_example.py" when testing "complete_cv_example.py"
+            special_strings (`list`, *optional*):
+                A list of strings to potentially remove before checking no differences are left. These should be
+                diffs that are file specific, such as different logging variations between files.
+        """
+        self.maxDiff = None
+        by_feature_path = os.path.abspath(os.path.join("examples", "by_feature"))
+        examples_path = os.path.abspath("examples")
+        for item in os.listdir(by_feature_path):
+            item_path = os.path.join(by_feature_path, item)
+            if os.path.isfile(item_path) and ".py" in item_path:
+                with self.subTest(
+                    tested_script=complete_file_name,
+                    feature_script=item,
+                    tested_section="main()" if parser_only else "training_function()",
+                ):
+                    diff = compare_against_test(
+                        os.path.join(examples_path, complete_file_name), item_path, parser_only, secondary_filename
+                    )
+                    diff = "\n".join(diff)
+                    if special_strings is not None:
+                        for string in special_strings:
+                            diff = diff.replace(string, "")
+                    self.assertEqual(diff, "")
+
+    def test_nlp_examples(self):
+        self.one_complete_example("complete_nlp_example.py", True)
+        self.one_complete_example("complete_nlp_example.py", False)
+
+    def test_cv_examples(self):
+        cv_path = os.path.abspath(os.path.join("examples", "cv_example.py"))
+        special_strings = [
+            " " * 16 + "{\n\n",
+            " " * 18 + '"accuracy": eval_metric["accuracy"],\n\n',
+            " " * 18 + '"f1": eval_metric["f1"],\n\n',
+            " " * 18 + '"train_loss": total_loss,\n\n',
+            " " * 18 + '"epoch": epoch,\n\n',
+            " " * 16 + "}\n",
+            " " * 8,
+        ]
+        self.one_complete_example("complete_cv_example.py", True, cv_path, special_strings)
+        self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
+
+
+class FeatureExamplesTests(unittest.TestCase):
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_checkpointing_by_epoch(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            testargs = f"""
+            checkpointing.py
+            --checkpointing_steps epoch
+            --output_dir {tmpdir}
+            """.split()
+            with mock.patch.object(sys, "argv", testargs):
+                checkpointing.main()
+                self.assertTrue(os.path.exists(os.path.join(tmpdir, "epoch_0")))
+
+    @mock.patch("checkpointing.get_dataloaders", mocked_dataloaders)
+    def test_checkpointing_by_steps(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            testargs = f"""
+            checkpointing.py
+            --checkpointing_steps 2
+            --output_dir {tmpdir}
+            """.split()
+            with mock.patch.object(sys, "argv", testargs):
+                checkpointing.main()
+                self.assertTrue(os.path.exists(os.path.join(tmpdir, "step_2")))
+
+    @mock.patch("tracking.get_dataloaders", mocked_dataloaders)
+    def test_tracking(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            testargs = f"""
+            tracking.py
+            --with_tracking
+            --logging_dir {tmpdir}
+            """.split()
+            with mock.patch.object(sys, "argv", testargs):
+                tracking.main()
+                self.assertTrue(os.path.exists(os.path.join(tmpdir, "tracking")))
diff --git a/tests/test_samples/MRPC/dev.csv b/tests/test_samples/MRPC/dev.csv
new file mode 100644
index 000000000..96beccda9
--- /dev/null
+++ b/tests/test_samples/MRPC/dev.csv
@@ -0,0 +1,7 @@
+label,sentence1,sentence2
+equivalent,He said the foodservice pie business doesn 't fit the company 's long-term growth strategy .,""" The foodservice pie business does not fit our long-term growth strategy ."
+not_equivalent,Magnarelli said Racicot hated the Iraqi regime and looked forward to using his long years of training in the war .,"His wife said he was "" 100 percent behind George Bush "" and looked forward to using his years of training in the war ."
+not_equivalent,"The dollar was at 116.92 yen against the yen , flat on the session , and at 1.2891 against the Swiss franc , also flat .","The dollar was at 116.78 yen JPY = , virtually flat on the session , and at 1.2871 against the Swiss franc CHF = , down 0.1 percent ."
+equivalent,The AFL-CIO is waiting until October to decide if it will endorse a candidate .,The AFL-CIO announced Wednesday that it will decide in October whether to endorse a candidate before the primaries .
+not_equivalent,No dates have been set for the civil or the criminal trial .,"No dates have been set for the criminal or civil cases , but Shanley has pleaded not guilty ."
+equivalent,Wal-Mart said it would check all of its million-plus domestic workers to ensure they were legally employed .,It has also said it would review all of its domestic employees more than 1 million to ensure they have legal status .
diff --git a/tests/test_samples/MRPC/train.csv b/tests/test_samples/MRPC/train.csv
new file mode 100644
index 000000000..96beccda9
--- /dev/null
+++ b/tests/test_samples/MRPC/train.csv
@@ -0,0 +1,7 @@
+label,sentence1,sentence2
+equivalent,He said the foodservice pie business doesn 't fit the company 's long-term growth strategy .,""" The foodservice pie business does not fit our long-term growth strategy ."
+not_equivalent,Magnarelli said Racicot hated the Iraqi regime and looked forward to using his long years of training in the war .,"His wife said he was "" 100 percent behind George Bush "" and looked forward to using his years of training in the war ."
+not_equivalent,"The dollar was at 116.92 yen against the yen , flat on the session , and at 1.2891 against the Swiss franc , also flat .","The dollar was at 116.78 yen JPY = , virtually flat on the session , and at 1.2871 against the Swiss franc CHF = , down 0.1 percent ."
+equivalent,The AFL-CIO is waiting until October to decide if it will endorse a candidate .,The AFL-CIO announced Wednesday that it will decide in October whether to endorse a candidate before the primaries .
+not_equivalent,No dates have been set for the civil or the criminal trial .,"No dates have been set for the criminal or civil cases , but Shanley has pleaded not guilty ."
+equivalent,Wal-Mart said it would check all of its million-plus domestic workers to ensure they were legally employed .,It has also said it would review all of its domestic employees more than 1 million to ensure they have legal status .


Review comments: body: Why not always do this?
diff_hunk: @@ -200,6 +200,8 @@ def training_function(config, args):
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
+            if isinstance(checkpointing_steps, int):
+                overall_step += 1
body: It's only needed if we save via checkpointing steps rather than epoch, so didn't want people to assume we *always* need to do that. 

(Which means a comment is needed!)
diff_hunk: @@ -200,6 +200,8 @@ def training_function(config, args):
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
+            if isinstance(checkpointing_steps, int):
+                overall_step += 1
body: Yeah I get the variable won't be used if we don't checkpoint with steps, but it doesn't hurt to always have it (and would save one line of code).
diff_hunk: @@ -200,6 +200,8 @@ def training_function(config, args):
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
+            if isinstance(checkpointing_steps, int):
+                overall_step += 1
body: ```suggestion
# Copyright 2022 The HuggingFace Team. All rights reserved.
```
diff_hunk: @@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
body: A quick intro here would be useful.
diff_hunk: @@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
body: ```suggestion
def get_train_func(lines: list):
```
If we use type annotations, better use `List[str]` here.
diff_hunk: @@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import os
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_train_func(lines: list):
body: Sounds like it could be refactored with the previous one and an extra arg?
diff_hunk: @@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import os
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_train_func(lines: list):
+    """
+    Finds the main training function from inside segmented source code.
+
+    Args:
+        lines (`list`):
+            Source code of a script
+    """
+    good_lines, found_start = [], False
+    for line in lines:
+        if not found_start and "def training_function" in line:
+            found_start = True
+            good_lines.append(line)
+            continue
+        if found_start:
+            if "def main" in line:
+                return good_lines
+            good_lines.append(line)
+
+
+def get_main_func(lines: list):
body: Those do not match the signature. Also this could be expanded a little bit, I'm not sure I fully understand.
diff_hunk: @@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import os
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_train_func(lines: list):
+    """
+    Finds the main training function from inside segmented source code.
+
+    Args:
+        lines (`list`):
+            Source code of a script
+    """
+    good_lines, found_start = [], False
+    for line in lines:
+        if not found_start and "def training_function" in line:
+            found_start = True
+            good_lines.append(line)
+            continue
+        if found_start:
+            if "def main" in line:
+                return good_lines
+            good_lines.append(line)
+
+
+def get_main_func(lines: list):
+    """
+    Finds the main function from inside segmented source code
+
+    Args:
+        lines (`list`):
+            Source code of a script
+    """
+    good_lines, found_start = [], False
+    for line in lines:
+        if not found_start and "def main" in line:
+            found_start = True
+            good_lines.append(line)
+            continue
+        if found_start:
+            if "if __name__" in line:
+                return good_lines
+            good_lines.append(line)
+
+
+def clean_lines(lines: list):
+    """
+    Filters `lines` and removes any entries that start with a comment ('#') or is just a newline ('\n')
+
+    Args:
+        lines (`list`):
+            Source code of a script
+    """
+    return [line for line in lines if not line.lstrip().startswith("#") and line != "\n"]
+
+
+def compare_against_test(base_filename: str, feature_filename: str, parser_only: bool, secondary_filename: str = None):
+    """
+    Checks whether the content aligned in `test_filename` is included inside of `full_filename`.
body: ```suggestion
# Copyright 2022 The HuggingFace Team. All rights reserved.
```
diff_hunk: @@ -0,0 +1,176 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
body: It seems we always use the method with both flags, should we just remove that arg and put the two tests inside?
diff_hunk: @@ -0,0 +1,176 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+import tempfile
+import unittest
+from unittest import mock
+
+from torch.utils.data import DataLoader
+
+from accelerate import DistributedType
+from accelerate.test_utils.examples import compare_against_test
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+SRC_DIRS = [os.path.abspath(os.path.join("examples", "by_feature"))]
+sys.path.extend(SRC_DIRS)
+
+if SRC_DIRS is not None:
+    import checkpointing
+    import tracking
+
+# DataLoaders built from `test_samples/MRPC` for quick testing
+# Should mock `{script_name}.get_dataloaders` via:
+# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
+
+
+def mocked_dataloaders(accelerator, batch_size: int = 16):
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
+    datasets = load_dataset("csv", data_files=data_files)
+    label_list = datasets["train"].unique("label")
+
+    label_to_id = {v: i for i, v in enumerate(label_list)}
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(
+            examples["sentence1"], examples["sentence2"], truncation=True, max_length=None, padding="max_length"
+        )
+        if "label" in examples:
+            outputs["labels"] = [label_to_id[l] for l in examples["label"]]
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["sentence1", "sentence2", "label"],
+    )
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=2)
+    eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1)
+
+    return train_dataloader, eval_dataloader
+
+
+class ExampleDifferenceTests(unittest.TestCase):
+    """
+    This TestCase checks that all of the `complete_*` scripts contain all of the
+    information found in the `by_feature` scripts, line for line.
+    """
+
+    def one_complete_example(
+        self, complete_file_name: str, parser_only: bool, secondary_filename: str = None, special_strings: list = None
+    ):
+        """
+        Tests a single `complete` example against all of the implemented `by_feature` scripts
+
+        Args:
+            complete_file_name (`str`):
+                The filename of a complete example
+            parser_only (`bool`):
+                Whether to look at the main training function, or the argument parser
+            secondary_filename (`str`, *optional*):
+                A potential secondary base file to strip all script information not relevant for checking,
+                such as "cv_example.py" when testing "complete_cv_example.py"
+            special_strings (`list`, *optional*):
+                A list of strings to potentially remove before checking no differences are left. These should be
+                diffs that are file specific, such as different logging variations between files.
+        """
+        self.maxDiff = None
+        by_feature_path = os.path.abspath(os.path.join("examples", "by_feature"))
+        examples_path = os.path.abspath("examples")
+        for item in os.listdir(by_feature_path):
+            item_path = os.path.join(by_feature_path, item)
+            if os.path.isfile(item_path) and ".py" in item_path:
+                with self.subTest(
+                    tested_script=complete_file_name,
+                    feature_script=item,
+                    tested_section="main()" if parser_only else "training_function()",
+                ):
+                    diff = compare_against_test(
+                        os.path.join(examples_path, complete_file_name), item_path, parser_only, secondary_filename
+                    )
+                    diff = "\n".join(diff)
+                    if special_strings is not None:
+                        for string in special_strings:
+                            diff = diff.replace(string, "")
+                    self.assertEqual(diff, "")
+
+    def test_nlp_examples(self):
+        self.one_complete_example("complete_nlp_example.py", True)
+        self.one_complete_example("complete_nlp_example.py", False)
+
body: I'm making a note for this inside of the documentation, but the reasoning for separation is it lets the test failure be more readable as to what section failed, rather than one complete error.
Notice the `tested_section` part
E.g.:

```bash
_________ ExampleDifferenceTests.test_cv_example (feature_script='tracking.py', tested_script='complete_cv_example.py', tested_section='training_function()') _________
```
vs:
```bash
_________ ExampleDifferenceTests.test_cv_example (feature_script='tracking.py', tested_script='complete_cv_example.py', tested_section='main()') _________
```

(This is a `pytest-subtest` hack)
diff_hunk: @@ -0,0 +1,176 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+import tempfile
+import unittest
+from unittest import mock
+
+from torch.utils.data import DataLoader
+
+from accelerate import DistributedType
+from accelerate.test_utils.examples import compare_against_test
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+
+SRC_DIRS = [os.path.abspath(os.path.join("examples", "by_feature"))]
+sys.path.extend(SRC_DIRS)
+
+if SRC_DIRS is not None:
+    import checkpointing
+    import tracking
+
+# DataLoaders built from `test_samples/MRPC` for quick testing
+# Should mock `{script_name}.get_dataloaders` via:
+# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)
+
+
+def mocked_dataloaders(accelerator, batch_size: int = 16):
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+    data_files = {"train": "tests/test_samples/MRPC/train.csv", "validation": "tests/test_samples/MRPC/dev.csv"}
+    datasets = load_dataset("csv", data_files=data_files)
+    label_list = datasets["train"].unique("label")
+
+    label_to_id = {v: i for i, v in enumerate(label_list)}
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(
+            examples["sentence1"], examples["sentence2"], truncation=True, max_length=None, padding="max_length"
+        )
+        if "label" in examples:
+            outputs["labels"] = [label_to_id[l] for l in examples["label"]]
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["sentence1", "sentence2", "label"],
+    )
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.TPU:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=2)
+    eval_dataloader = DataLoader(tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=1)
+
+    return train_dataloader, eval_dataloader
+
+
+class ExampleDifferenceTests(unittest.TestCase):
+    """
+    This TestCase checks that all of the `complete_*` scripts contain all of the
+    information found in the `by_feature` scripts, line for line.
+    """
+
+    def one_complete_example(
+        self, complete_file_name: str, parser_only: bool, secondary_filename: str = None, special_strings: list = None
+    ):
+        """
+        Tests a single `complete` example against all of the implemented `by_feature` scripts
+
+        Args:
+            complete_file_name (`str`):
+                The filename of a complete example
+            parser_only (`bool`):
+                Whether to look at the main training function, or the argument parser
+            secondary_filename (`str`, *optional*):
+                A potential secondary base file to strip all script information not relevant for checking,
+                such as "cv_example.py" when testing "complete_cv_example.py"
+            special_strings (`list`, *optional*):
+                A list of strings to potentially remove before checking no differences are left. These should be
+                diffs that are file specific, such as different logging variations between files.
+        """
+        self.maxDiff = None
+        by_feature_path = os.path.abspath(os.path.join("examples", "by_feature"))
+        examples_path = os.path.abspath("examples")
+        for item in os.listdir(by_feature_path):
+            item_path = os.path.join(by_feature_path, item)
+            if os.path.isfile(item_path) and ".py" in item_path:
+                with self.subTest(
+                    tested_script=complete_file_name,
+                    feature_script=item,
+                    tested_section="main()" if parser_only else "training_function()",
+                ):
+                    diff = compare_against_test(
+                        os.path.join(examples_path, complete_file_name), item_path, parser_only, secondary_filename
+                    )
+                    diff = "\n".join(diff)
+                    if special_strings is not None:
+                        for string in special_strings:
+                            diff = diff.replace(string, "")
+                    self.assertEqual(diff, "")
+
+    def test_nlp_examples(self):
+        self.one_complete_example("complete_nlp_example.py", True)
+        self.one_complete_example("complete_nlp_example.py", False)
+
body: Should go inside triple quotes.
diff_hunk: @@ -0,0 +1,143 @@
+#!/usr/bin/env python
+
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# A collection of utilities for comparing `examples/complete_*_example.py` scripts with the capabilities inside of each
+# `examples/by_feature` example. `compare_against_test` is the main function that should be used when testing,
+# while the others are used to either get the code that matters, or to preprocess them (such as stripping comments)

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
body: I've come up with a solution to make sure that also the main `complete_` examples will raise an error if they differ from the bits inside of `by_feature`. Now eventually it would be good to have an exclusion list of files for when we *wouldn't* expect to see some feature added (such as say cross validation). But we'll cross that bridge when we get there. In the meantime here is the report generated back to the user from pytest:

(Note: `pytest-subtests` was added as a dep to make sure we could use `TestCase.subTest` and make this code clean and efficient)
```bash
======================================================================= short test summary info =======================================================================
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_cv_example_body - AssertionError: 14 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_cv_example_body - AssertionError: 27 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_cv_example_parser - AssertionError: 1 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_cv_example_parser - AssertionError: 4 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_nlp_example_parser - AssertionError: 1 != 0
SUBFAIL tests/test_examples.py::ExampleDifferenceTests::test_complete_nlp_example_parser - AssertionError: 1 != 0
=============================================================== 6 failed, 4 passed, 1 warning in 1.64s ================================================================
```

{{ Removed full trace, see the following message for an example }} 
body: Slightly tweaked how they look, I've now included it posting the source code diffs. I believe this is a must as it will tell the user exactly what parts were missing from the full example:

```bash
============================================================================== FAILURES ===============================================================================
_________ ExampleDifferenceTests.test_cv_example (feature_script='tracking.py', tested_script='complete_cv_example.py', tested_section='training_function()') _________

self = <test_examples.ExampleDifferenceTests testMethod=test_cv_example>, complete_file_name = 'complete_cv_example.py', parser_only = False

    def one_complete_example(self, complete_file_name: str, parser_only: bool):
        """
        Tests a single `complete` example against all of the implemented `by_feature` scripts
    
        Args:
            complete_file_name (`str`):
                The filename of a complete example
            parser_only (`bool`):
                Whether to look at the main training function, or the argument parser
        """
        self.maxDiff = None
        by_feature_path = os.path.abspath(os.path.join("examples", "by_feature"))
        examples_path = os.path.abspath("examples")
        for item in os.listdir(by_feature_path):
            item_path = os.path.join(by_feature_path, item)
            if os.path.isfile(item_path) and ".py" in item_path:
                with self.subTest(tested_script=complete_file_name, feature_script=item, tested_section="main()" if parser_only else "training_function()"):
                    diff = compare_against_test(
                        os.path.join(examples_path, "nlp_example.py"),
                        os.path.join(examples_path, complete_file_name),
                        item_path,
                        parser_only
                    )
>                   self.assertEqual('\n'.join(diff), '')
E                   AssertionError: '        accelerator = Accelerator(cpu=arg[693 chars]()\n' != ''
E                   -         accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision, log_with="all")
E                   - 
E                   -         accelerator = Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)
E                   - 
E                   -         accelerator.init_trackers("nlp_example", config)
E                   - 
E                   -             predictions, references = accelerator.gather((predictions, batch["labels"]))
E                   - 
E                   -                 predictions=predictions,
E                   - 
E                   -                 references=references,
E                   - 
E                   -             accelerator.log(
E                   - 
E                   -                 {
E                   - 
E                   -                     "accuracy": eval_metric["accuracy"],
E                   - 
E                   -                     "f1": eval_metric["f1"],
E                   - 
E                   -                     "train_loss": total_loss,
E                   - 
E                   -                     "epoch": epoch,
E                   - 
E                   -                 }
E                   - 
E                   -         accelerator.end_training()

tests/test_examples.py:107: AssertionError
```