diff --git a/stable-diffusion-finetuning/pokemon_finetune.ipynb b/stable-diffusion-finetuning/pokemon_finetune.ipynb
index 79aa765..5d6b095 100644
--- a/stable-diffusion-finetuning/pokemon_finetune.ipynb
+++ b/stable-diffusion-finetuning/pokemon_finetune.ipynb
@@ -1,6 +1,7 @@
 {
   "cells": [
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ipSgY-L3z8kG"
@@ -59,6 +60,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "5V1uUVjQz8kI"
@@ -94,6 +96,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "y8VDyuQxz8kJ"
@@ -118,9 +121,21 @@
         "BATCH_SIZE = 4\n",
         "N_GPUS = 2\n",
         "ACCUMULATE_BATCHES = 1\n",
+        "GPU_LIST = \",\".join((str(x) for x in range(N_GPUS))) + \",\"\n",
+        "print(f\"Using GPUs: {GPU_LIST}\")\n",
         "\n",
-        "gpu_list = \",\".join((str(x) for x in range(N_GPUS))) + \",\"\n",
-        "print(f\"Using GPUs: {gpu_list}\")"
+        "import os\n",
+        "os.environ[\"BATCH_SIZE\"] = str(BATCH_SIZE) \n",
+        "os.environ[\"N_GPUS\"] = str(N_GPUS) \n",
+        "os.environ[\"ACCUMULATE_BATCHES\"] = str(ACCUMULATE_BATCHES)\n",
+        "os.environ[\"GPU_LIST\"] = GPU_LIST\n",
+        "os.environ[\"CKPT_PATH\"] = ckpt_path\n",
+        "\n",
+        "!echo \"$BATCH_SIZE\"\n",
+        "!echo \"$N_GPUS\"\n",
+        "!echo \"$ACCUMULATE_BATCHES\"\n",
+        "!echo \"$GPU_LIST\"\n",
+        "!echo \"$CKPT_PATH\""
       ]
     },
     {
@@ -135,14 +150,14 @@
         "!(python main.py \\\n",
         "    -t \\\n",
         "    --base configs/stable-diffusion/pokemon.yaml \\\n",
-        "    --gpus \"$gpu_list\" \\\n",
+        "    --gpus \"$GPU_LIST\" \\\n",
         "    --scale_lr False \\\n",
         "    --num_nodes 1 \\\n",
         "    --check_val_every_n_epoch 10 \\\n",
-        "    --finetune_from \"$ckpt_path\" \\\n",
+        "    --finetune_from \"$CKPT_PATH\" \\\n",
         "    data.params.batch_size=\"$BATCH_SIZE\" \\\n",
         "    lightning.trainer.accumulate_grad_batches=\"$ACCUMULATE_BATCHES\" \\\n",
-        "    data.params.validation.params.n_gpus=\"$NUM_GPUS\" \\\n",
+        "    data.params.validation.params.n_gpus=\"$N_GPUS\" \\\n",
         ")"
       ]
     },