diff --git a/ehrapy/plot/_survival_analysis.py b/ehrapy/plot/_survival_analysis.py
index 7a7eff72..2045f93e 100644
--- a/ehrapy/plot/_survival_analysis.py
+++ b/ehrapy/plot/_survival_analysis.py
@@ -269,7 +269,7 @@ def kaplan_meier(
         color = [None] * len(kmfs)
 
     fig = plt.figure(constrained_layout=True, figsize=figsize)
-    spec = fig.add_gridspec(2, 1)
+    spec = fig.add_gridspec(2, 1) if display_survival_statistics else fig.add_gridspec(1, 1)
     ax = plt.subplot(spec[0, 0])
 
     for i, kmf in enumerate(kmfs):
@@ -301,7 +301,6 @@ def kaplan_meier(
     if title:
         ax.set_title(title)
 
-    # Prepare data for the table
     if display_survival_statistics:
         xticks = [x for x in ax.get_xticks() if x >= 0]
         xticks_space = xticks[1] - xticks[0]
diff --git a/tests/_scripts/kaplain_meier_create_expected_plots.ipynb b/tests/_scripts/kaplain_meier_create_expected_plots.ipynb
index b5bff552..21cf88ab 100644
--- a/tests/_scripts/kaplain_meier_create_expected_plots.ipynb
+++ b/tests/_scripts/kaplain_meier_create_expected_plots.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -27,7 +27,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -58,7 +58,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -86,7 +86,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
diff --git a/tests/plot/_images/kaplan_meier_expected.png b/tests/plot/_images/kaplan_meier_expected.png
index cee1a8d0..8b9c619d 100644
Binary files a/tests/plot/_images/kaplan_meier_expected.png and b/tests/plot/_images/kaplan_meier_expected.png differ