diff --git a/base/scripts/run-notebook.py b/base/scripts/run-notebook.py
index d78c646..7606cad 100755
--- a/base/scripts/run-notebook.py
+++ b/base/scripts/run-notebook.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import argparse
 import sys
 from pathlib import Path
 import logging
@@ -9,6 +10,16 @@
 from papermill.utils import merge_kwargs, remove_args
 from nbconvert import HTMLExporter
 
+
+parser = argparse.ArgumentParser(
+    prog='run-notebook.py',
+    description='Run a notebook and converts it to HTML while it is running')
+
+parser.add_argument('notebook_path', help="The path to the .ipynb notebook file")
+parser.add_argument('notebook_output_path', help="The path where the .ipynb notebook file should be written")
+parser.add_argument('-c', '--cwd', help="Change the working directory")
+args = parser.parse_args()
+
 logging.basicConfig(level=logging.INFO, format="%(message)s")
 logger.setLevel(logging.INFO)
 
@@ -68,8 +79,9 @@ def execute_notebook(
 
 # Execute the notebook with papermill as usual
 papermill.execute_notebook(
-    sys.argv[1],
-    sys.argv[2],
+    args.notebook_path,
+    args.notebook_output_path,
+    cwd=args.cwd,
     log_output=True,
     engine_name="htmlengine",
 )
\ No newline at end of file
diff --git a/base/scripts/run-notebook.sh b/base/scripts/run-notebook.sh
index 593ff36..617871a 100755
--- a/base/scripts/run-notebook.sh
+++ b/base/scripts/run-notebook.sh
@@ -5,21 +5,18 @@ if [[ -z "$NOTEBOOK" ]]; then
   for nb in /content/*.ipynb; do
     [ -e "$nb" ] || continue
     echo "Executing notebook $nb"
-    run-notebook.py "$nb" "/content/artifacts/$(basename ${nb})"
-    # jupyter nbconvert --to html /content/artifacts/$(basename "$nb")
+    run-notebook.py "$nb" "/content/artifacts/$(basename ${nb})" --cwd "$(dirname ${nb})"
   done
   for nb in /content/src/*.ipynb; do
     [ -e "$nb" ] || continue
     mkdir -p /content/artifacts/src
     echo "Executing notebook $nb"
-    run-notebook.py "$nb" "/content/artifacts/src/$(basename ${nb})"
-    # jupyter nbconvert --to html /content/artifacts//src/$(basename "$nb")
+    run-notebook.py "$nb" "/content/artifacts/src/$(basename ${nb})" --cwd "$(dirname ${nb})"
   done
 else
   for nb in $NOTEBOOK; do
     [ -e "$nb" ] || continue
     echo "Executing notebook $nb"
-    run-notebook.py "$nb" "/content/artifacts/$(basename ${nb})"
-    # jupyter nbconvert --to html /content/artifacts/$(basename "$nb")
+    run-notebook.py "$nb" "/content/artifacts/$(basename ${nb})" --cwd "$(dirname ${nb})"
   done
 fi
\ No newline at end of file
diff --git a/model-loader-huggingface/Dockerfile b/model-loader-huggingface/Dockerfile
index 85555a2..ae72605 100644
--- a/model-loader-huggingface/Dockerfile
+++ b/model-loader-huggingface/Dockerfile
@@ -6,7 +6,4 @@ RUN mkdir -p /content/model
 COPY requirements.txt requirements.txt
 RUN pip3 install --no-cache-dir -r requirements.txt
 
-COPY load.sh /usr/local/bin/
 COPY src/* src/
-
-CMD load.sh
diff --git a/model-loader-huggingface/load.sh b/model-loader-huggingface/load.sh
deleted file mode 100755
index f842946..0000000
--- a/model-loader-huggingface/load.sh
+++ /dev/null
@@ -1,5 +0,0 @@
-#!/usr/bin/env sh
-
-set -xe
-
-jupyter nbconvert --debug --to notebook --execute /content/src/load.ipynb --output /content/logs/load.ipynb
diff --git a/model-trainer-huggingface/Dockerfile b/model-trainer-huggingface/Dockerfile
index 13b615b..e8fe948 100644
--- a/model-trainer-huggingface/Dockerfile
+++ b/model-trainer-huggingface/Dockerfile
@@ -7,7 +7,4 @@ WORKDIR /content
 COPY requirements.txt /content/requirements.txt
 RUN pip install --no-cache-dir -r /content/requirements.txt
 
-COPY scripts/* ./scripts/
 COPY src/* ./src/
-
-CMD train.sh
diff --git a/model-trainer-huggingface/scripts/train.sh b/model-trainer-huggingface/scripts/train.sh
deleted file mode 100755
index 30186b1..0000000
--- a/model-trainer-huggingface/scripts/train.sh
+++ /dev/null
@@ -1,6 +0,0 @@
-#!/usr/bin/env sh
-
-set -xe
-
-export HF_HUB_DISABLE_PROGRESS_BARS=1
-jupyter nbconvert --debug --to notebook --execute /content/src/train.ipynb --output /content/logs/train.ipynb