diff --git a/03_COSDD/exercise.ipynb b/03_COSDD/exercise.ipynb
index ba95cdc..b1dc8d5 100644
--- a/03_COSDD/exercise.ipynb
+++ b/03_COSDD/exercise.ipynb
@@ -30,7 +30,6 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "import logging\n",
     "\n",
     "import torch\n",
     "import tifffile\n",
@@ -307,8 +306,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "real_batch_size = 16\n",
-    "n_grad_batches = 1\n",
+    "real_batch_size = 4\n",
+    "n_grad_batches = 4\n",
     "print(f\"Effective batch size: {real_batch_size * n_grad_batches}\")\n",
     "crop_size = (256, 256)\n",
     "train_split = 0.9\n",
@@ -408,9 +407,9 @@
    "outputs": [],
    "source": [
     "dimensions = ... ### Insert a value here\n",
-    "s_code_channels = 16\n",
+    "s_code_channels = 32\n",
     "\n",
-    "n_layers = 4\n",
+    "n_layers = 6\n",
     "z_dims = [s_code_channels // 2] * n_layers\n",
     "downsampling = [1] * n_layers\n",
     "lvae = LadderVAE(\n",
@@ -428,8 +427,8 @@
     "    s_code_channels=s_code_channels,\n",
     "    kernel_size=5,\n",
     "    noise_direction=...  ### Insert a value here\n",
-    "    n_filters=16,\n",
-    "    n_layers=3,\n",
+    "    n_filters=32,\n",
+    "    n_layers=4,\n",
     "    n_gaussians=4,\n",
     "    dimensions=dimensions,\n",
     ")\n",
@@ -462,7 +461,7 @@
     "    data_mean=low_snr.mean(),\n",
     "    data_std=low_snr.std(),\n",
     "    n_grad_batches=n_grad_batches,\n",
-    "    checkpointed=False,\n",
+    "    checkpointed=True,\n",
     ")"
    ]
   },
@@ -579,7 +578,6 @@
     "    max_time=max_time,  # Remove this time limit to train the model fully\n",
     "    log_every_n_steps=len(train_set) // (n_grad_batches * real_batch_size),\n",
     "    callbacks=[EarlyStopping(patience=patience, monitor=\"val/elbo\")],\n",
-    "    precision=\"bf16-mixed\",\n",
     ")"
    ]
   },
@@ -613,22 +611,12 @@
     "# Exercise 2. Inference with COSDD"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "logger = logging.getLogger('pytorch_lightning')\n",
-    "logger.setLevel(logging.WARNING)"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### 2.1. Load test data\n",
-    "The images that we want to denoise are loaded here. These are the same that we used for training, but we'll only load 2 to speed up inference."
+    "The images that we want to denoise are loaded here. These are the same that we used for training, but we'll only load 10 to speed up inference."
    ]
   },
   {
@@ -638,7 +626,7 @@
    "outputs": [],
    "source": [
     "lowsnr_path = \"./../data/mito-confocal-lowsnr.tif\"\n",
-    "n_test_images = 2\n",
+    "n_test_images = 10\n",
     "# load the data\n",
     "test_set = tifffile.imread(lowsnr_path)\n",
     "test_set = test_set[:n_test_images, np.newaxis]\n",
@@ -724,7 +712,6 @@
     "    enable_progress_bar=False,\n",
     "    enable_checkpointing=False,\n",
     "    logger=False,\n",
-    "    precision=\"bf16-mixed\",\n",
     ")"
    ]
   },
@@ -747,7 +734,6 @@
     "    enable_progress_bar=False,\n",
     "    enable_checkpointing=False,\n",
     "    logger=False,\n",
-    "    precision=\"bf16-mixed\",\n",
     ")"
    ]
   },
diff --git a/04_DenoiSplit/exercise.ipynb b/04_DenoiSplit/exercise.ipynb
index e269f84..b84a78f 100644
--- a/04_DenoiSplit/exercise.ipynb
+++ b/04_DenoiSplit/exercise.ipynb
@@ -796,7 +796,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ccb2694d",
+   "id": "337b207c",
    "metadata": {},
    "source": [
     "<hr style=\"height:2px;\"><div class=\"alert alert-block alert-success\"><h1>End of the exercise</h1>\n",