diff --git a/colabcode/code.py b/colabcode/code.py
index 6432b83..5362158 100644
--- a/colabcode/code.py
+++ b/colabcode/code.py
@@ -14,9 +14,10 @@
 
 
 class ColabCode:
-    def __init__(self, port=10000, password=None, mount_drive=False):
+    def __init__(self, port=10000, password=None, user_data_dir=None, mount_drive=False):
         self.port = port
         self.password = password
+        self.user_data_dir = user_data_dir
         self._mount = mount_drive
         self._install_code()
         self._install_extensions()
@@ -45,8 +46,12 @@ def _run_code(self):
         os.system(f"fuser -n tcp -k {self.port}")
         if self._mount and colab_env:
             drive.mount("/content/drive")
-        if self.password:
+        if self.password and self.user_data_dir:
+            code_cmd = f"PASSWORD={self.password} code-server --port {self.port} --user-data-dir {self.user_data_dir} --disable-telemetry"
+        elif self.password:
             code_cmd = f"PASSWORD={self.password} code-server --port {self.port} --disable-telemetry"
+        elif self.user_data_dir:
+            code_cmd = f"code-server --port {self.port} --auth none --user-data-dir {self.user_data_dir} --disable-telemetry"
         else:
             code_cmd = f"code-server --port {self.port} --auth none --disable-telemetry"
         with subprocess.Popen(