diff --git a/docs/source/use_with_jax.mdx b/docs/source/use_with_jax.mdx
index 7c601e7be37..dc73a8df778 100644
--- a/docs/source/use_with_jax.mdx
+++ b/docs/source/use_with_jax.mdx
@@ -77,7 +77,7 @@ True
 Note that if the `device` argument is not provided to `with_format` then it will use the default
 device which is `jax.devices()[0]`.
 
-## N-dimensional arrays
+### N-dimensional arrays
 
 If your dataset consists of N-dimensional arrays, you will see that by default they are considered as the same tensor if the shape is fixed:
 
@@ -120,7 +120,7 @@ To avoid this, you must explicitly use the [`Array`] feature type and specify th
          [7, 8]]], dtype=int32)}
 ```
 
-## Other feature types
+### Other feature types
 
 [`ClassLabel`] data is properly converted to arrays:
 
diff --git a/docs/source/use_with_pytorch.mdx b/docs/source/use_with_pytorch.mdx
index dae29557694..7f78d8de05c 100644
--- a/docs/source/use_with_pytorch.mdx
+++ b/docs/source/use_with_pytorch.mdx
@@ -38,7 +38,7 @@ To load the data as tensors on a GPU, specify the `device` argument:
 {'data': tensor([1, 2], device='cuda:0')}
 ```
 
-## N-dimensional arrays
+### N-dimensional arrays
 
 If your dataset consists of N-dimensional arrays, you will see that by default they are considered as the same tensor if the shape is fixed:
 
@@ -82,7 +82,7 @@ To avoid this, you must explicitly use the [`Array`] feature type and specify th
 ```
 
 
-## Other feature types
+### Other feature types
 
 [`ClassLabel`] data are properly converted to tensors:
 
@@ -223,6 +223,8 @@ If the dataset is split in several shards (i.e. if the dataset consists of multi
 
 In this case each worker is given a subset of the list of shards to stream from.
 
+### Checkpoint and resume
+
 If you need a DataLoader that you can checkpoint and resume in the middle of training, you can use the `StatefulDataLoader` from [torchdata](https://github.com/pytorch/data):
 
 ```py
diff --git a/docs/source/use_with_tensorflow.mdx b/docs/source/use_with_tensorflow.mdx
index ef16945024b..14e0ac2b33c 100644
--- a/docs/source/use_with_tensorflow.mdx
+++ b/docs/source/use_with_tensorflow.mdx
@@ -39,7 +39,7 @@ array([[1, 2],
        [3, 4]])>}
 ```
 
-## N-dimensional arrays
+### N-dimensional arrays
 
 If your dataset consists of N-dimensional arrays, you will see that by default they are considered as the same tensor if the shape is fixed:
 
@@ -88,7 +88,7 @@ To avoid this, you must explicitly use the [`Array`] feature type and specify th
 ```
 
 
-## Other feature types
+### Other feature types
 
 [`ClassLabel`] data are properly converted to tensors: