diff --git a/docs/guides/training_101/index.md b/docs/guides/training_101/index.md
index 0db5ee9..15aa422 100644
--- a/docs/guides/training_101/index.md
+++ b/docs/guides/training_101/index.md
@@ -152,7 +152,7 @@ from PIL import Image
 from refiners.fluxion.utils import image_to_tensor
 
 def generate_mask(size: int, seed: int | None = None) -> Generator[torch.Tensor, None, None]:
-    """Generate a tensor of a grayscale mask of size `size` using random rectangles."""
+    """Generate a tensor of a binary mask of size `size` using random rectangles."""
     if seed is None:
         seed = random.randint(0, 2**32 - 1)
     random.seed(seed)
@@ -186,10 +186,9 @@ mask = next(generate_mask(64, seed=42))
 tensor_to_image(mask).save("mask.png")
 ```
 
-Here are a few examples of generated images:
+Here are a two examples of generated masks:
 ![alt text](sample-0.png)
 ![alt text](sample-1.png)
-![alt text](sample-2.png)
 
 ## Trainer
 
@@ -445,6 +444,10 @@ class AutoencoderTrainer(Trainer[AutoencoderConfig, Batch]):
 
 ![alt text](loss-logging.png)
 
+## Wrap up
+
+You can train this toy model using the code below:
+
 ??? complete end-to-end code "Expand to see the full code."
 
     ```py
diff --git a/docs/guides/training_101/sample-0.png b/docs/guides/training_101/sample-0.png
index 0a2411d..e9f1c5f 100644
Binary files a/docs/guides/training_101/sample-0.png and b/docs/guides/training_101/sample-0.png differ
diff --git a/docs/guides/training_101/sample-1.png b/docs/guides/training_101/sample-1.png
index 00ec019..b44a582 100644
Binary files a/docs/guides/training_101/sample-1.png and b/docs/guides/training_101/sample-1.png differ
diff --git a/docs/guides/training_101/sample-2.png b/docs/guides/training_101/sample-2.png
deleted file mode 100644
index babbac6..0000000
Binary files a/docs/guides/training_101/sample-2.png and /dev/null differ