diff --git a/.gitignore b/.gitignore
index 8785181..b33ccb9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -17,3 +17,4 @@ serket/experimental/test_conv.py
 docs/_build
 fft_intro.ipynb
 lenna.png
+docs/notebooks/transformer_weights.pickle
diff --git a/docs/notebooks/train_convlstm.ipynb b/docs/notebooks/train_convlstm.ipynb
index 6cef29f..f0a3aa5 100644
--- a/docs/notebooks/train_convlstm.ipynb
+++ b/docs/notebooks/train_convlstm.ipynb
@@ -6,7 +6,7 @@
         "id": "PJyyxjLA2A62"
       },
       "source": [
-        "# ![](https://img.shields.io/badge/time%20series-FFF2DB) Train Convolutional `LSTM`\n",
+        "# Train Convolutional `LSTM`\n",
         "\n",
         "\n",
         "In this notebook, a simple next frame video prediction using convolutional `LSTM` is demonstrated."
diff --git a/docs/notebooks/train_fourier_features_network.ipynb b/docs/notebooks/train_fourier_features_network.ipynb
index 9bc4dae..8b4e159 100644
--- a/docs/notebooks/train_fourier_features_network.ipynb
+++ b/docs/notebooks/train_fourier_features_network.ipynb
@@ -5,7 +5,7 @@
    "id": "b2b68cc9-825c-490f-9476-7b43acec0eaa",
    "metadata": {},
    "source": [
-    "# ![](https://img.shields.io/badge/vision-A53662) Train fourier features network\n",
+    "# Train fourier features network\n",
     "\n",
     "_Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains_\n",
     "\n",
diff --git a/docs/notebooks/train_mnist.ipynb b/docs/notebooks/train_mnist.ipynb
index 3e6ea72..0035668 100644
--- a/docs/notebooks/train_mnist.ipynb
+++ b/docs/notebooks/train_mnist.ipynb
@@ -6,7 +6,7 @@
         "id": "k0i-yEycZkF-"
       },
       "source": [
-        "# ![](https://img.shields.io/badge/vision-A53662) Train `MNIST`"
+        "# Train `MNIST`"
       ]
     },
     {
diff --git a/docs/notebooks/train_pinn_burgers.ipynb b/docs/notebooks/train_pinn_burgers.ipynb
index 72b827d..be4e806 100644
--- a/docs/notebooks/train_pinn_burgers.ipynb
+++ b/docs/notebooks/train_pinn_burgers.ipynb
@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# ![](https://img.shields.io/badge/science-D1B8C0) Train `PINN`\n",
+    "# Train `PINN`\n",
     "\n",
     "In this example, [physics informed neural network (PINN)](https://maziarraissi.github.io/PINNs/) technique is used to train [burgers equation](https://en.wikipedia.org/wiki/Burgers%27_equation).\n",
     "\n",
diff --git a/docs/notebooks/train_transformer.ipynb b/docs/notebooks/train_transformer.ipynb
index f7efacc..8e2654b 100644
--- a/docs/notebooks/train_transformer.ipynb
+++ b/docs/notebooks/train_transformer.ipynb
@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# ![](https://img.shields.io/badge/language-FFAC2A)   Train transformer\n",
+    "# Train transformer\n",
     "\n",
     "In this guide,  a simple transformer model is trained from scratch to perform machine translation as detailed in [Attention is all you need](https://arxiv.org/abs/1706.03762) from Arabic to English. But you are free to choose any language pair from the configuration."
    ]
diff --git a/docs/notebooks/train_unet.ipynb b/docs/notebooks/train_unet.ipynb
index 23bc562..461554c 100644
--- a/docs/notebooks/train_unet.ipynb
+++ b/docs/notebooks/train_unet.ipynb
@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# ![](https://img.shields.io/badge/vision-A53662) Train `UNet` segmenter\n",
+    "# Train `UNet` segmenter\n",
     "\n",
     "\n",
     "\n",
diff --git a/docs/training_guides.rst b/docs/training_guides.rst
index 6c3802c..6d06ac3 100644
--- a/docs/training_guides.rst
+++ b/docs/training_guides.rst
@@ -11,4 +11,5 @@
     notebooks/train_pinn_burgers
     notebooks/train_fourier_features_network
     notebooks/train_convlstm
-    notebooks/train_unet
\ No newline at end of file
+    notebooks/train_unet
+    notebooks/train_transformer
\ No newline at end of file
diff --git a/tests/test_clustering.py b/tests/test_clustering.py
index 1d15c09..a2ead62 100644
--- a/tests/test_clustering.py
+++ b/tests/test_clustering.py
@@ -18,13 +18,13 @@
 import jax.numpy as jnp
 import jax.random as jr
 import numpy.testing as npt
-
+import pytest
 import serket as sk
 
 # Suppress FutureWarning
 warnings.simplefilter(action="ignore", category=FutureWarning)
 
-
+@pytest.mark.skip(reason="flaky test")
 def test_kmeans():
     from sklearn.cluster import KMeans