diff --git a/llm_sandbox/docker.py b/llm_sandbox/docker.py
index 80d73ef..1002c7f 100644
--- a/llm_sandbox/docker.py
+++ b/llm_sandbox/docker.py
@@ -31,6 +31,7 @@ def __init__(
         dockerfile: Optional[str] = None,
         lang: str = SupportedLanguage.PYTHON,
         keep_template: bool = False,
+        commit_container: bool = True,
         verbose: bool = False,
         mounts: Optional[list[Mount]] = None,
     ):
@@ -41,6 +42,7 @@ def __init__(
         :param dockerfile: Path to the Dockerfile, if image is not provided
         :param lang: Language of the code
         :param keep_template: if True, the image and container will not be removed after the session ends
+        :param commit_container: if True, the Docker container will be commited after the session ends
         :param verbose: if True, print messages
         """
         super().__init__(lang, verbose)
@@ -71,6 +73,7 @@ def __init__(
         self.container: Optional[Container] = None
         self.path = None
         self.keep_template = keep_template
+        self.commit_container = commit_container
         self.is_create_template: bool = False
         self.verbose = verbose
         self.mounts = mounts
@@ -114,7 +117,7 @@ def open(self):
 
     def close(self):
         if self.container:
-            if isinstance(self.image, Image):
+            if self.commit_container and isinstance(self.image, Image):
                 self.container.commit(self.image.tags[-1])
 
             self.container.remove(force=True)
diff --git a/llm_sandbox/session.py b/llm_sandbox/session.py
index 71bf4b4..5e56dfa 100644
--- a/llm_sandbox/session.py
+++ b/llm_sandbox/session.py
@@ -14,6 +14,7 @@ def __new__(
         dockerfile: Optional[str] = None,
         lang: str = SupportedLanguage.PYTHON,
         keep_template: bool = False,
+        commit_container: bool = True,
         verbose: bool = False,
         use_kubernetes: bool = False,
         kube_namespace: Optional[str] = "default",
@@ -25,6 +26,7 @@ def __new__(
         :param dockerfile: Path to the Dockerfile, if image is not provided
         :param lang: Language of the code
         :param keep_template: if True, the image and container will not be removed after the session ends
+        :param commit_container: if True, the Docker container will be commited after the session ends
         :param verbose: if True, print messages (default is True)
         :param use_kubernetes: if True, use Kubernetes instead of Docker (default is False)
         :param kube_namespace: Kubernetes namespace to use (only if 'use_kubernetes' is True), default is 'default'
@@ -46,5 +48,6 @@ def __new__(
             dockerfile=dockerfile,
             lang=lang,
             keep_template=keep_template,
+            commit_container=commit_container,
             verbose=verbose,
         )