diff --git a/scripts/speedrun_modal.py b/scripts/speedrun_modal.py
index 0e579b7..4e634f4 100644
--- a/scripts/speedrun_modal.py
+++ b/scripts/speedrun_modal.py
@@ -15,19 +15,24 @@ from pathlib import Path
 # Define the Modal app
 app = modal.App("nanochat-train")
 path_repo = "/root/nanochat"
+path_nanochat_base_dir="/root/.cache/nanochat"
+volume_nanochat_base_dir="nanochat-base-dir"
 
 # Create the image with all necessary dependencies
 image = (
     modal.Image.debian_slim(python_version="3.12")
-    .apt_install("curl")
-    .add_local_dir('.', path_repo, ignore=modal.FilePatternMatcher.from_file(f".gitignore"))
+    .apt_install("curl", "unzip")
+    .add_local_dir(".", path_repo, ignore=modal.FilePatternMatcher.from_file(f".gitignore"))
 )
 
 # Define the main function
+volume = modal.Volume.from_name(volume_nanochat_base_dir, create_if_missing=True)
+
 n_nodes = 1
 @app.function(
     image=image,
     gpu="H100:8",
+    volumes={path_nanochat_base_dir: volume},
     timeout=6 * 60 * 60,  # set 6 hours timeout since modal may take long time to request gpus
 )
 @modal.experimental.clustered(size=n_nodes)
@@ -37,5 +42,8 @@ def nanochat_train():
     # change working directory to repo
     os.chdir(path_repo)
 
-    command = 'chmod +x ./speedrun.sh && ./speedrun.sh'
-    subprocess.run(command, shell=True, check=True)
\ No newline at end of file
+    try:
+        command = "chmod +x ./speedrun.sh && ./speedrun.sh"
+        subprocess.run(command, shell=True, check=True)
+    finally:
+        volume.commit()
\ No newline at end of file