diff --git a/ai_edge_torch/__init__.py b/ai_edge_torch/__init__.py
index 99580428..35f149a5 100644
--- a/ai_edge_torch/__init__.py
+++ b/ai_edge_torch/__init__.py
@@ -17,6 +17,7 @@
 from .convert.converter import signature
 from .convert.to_channel_last_io import to_channel_last_io
 from .model import Model
+from .version import __version__
 
 
 def load(path: str) -> Model:
diff --git a/ai_edge_torch/version.py b/ai_edge_torch/version.py
new file mode 100644
index 00000000..564389aa
--- /dev/null
+++ b/ai_edge_torch/version.py
@@ -0,0 +1,16 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+__version__ = "0.2.0"
diff --git a/setup.py b/setup.py
index f7108ba6..709b860d 100644
--- a/setup.py
+++ b/setup.py
@@ -15,6 +15,7 @@
 
 import os
 import pathlib
+import re
 
 from setuptools import find_packages
 from setuptools import setup
@@ -33,11 +34,19 @@
 """.lstrip()
 
 name = "ai-edge-torch"
-version = "0.2.0"
+# TODO(b/357076369): move version updating logics to version.py
+version_py = here / "ai_edge_torch" / "version.py"
+version_regex = "__version__\s*=\s*(\"|')(?P<version>[^\"']+)(\"|')"
+version = re.search(version_regex, version_py.read_text()).group("version")
+
 if nightly_release_date := os.environ.get("NIGHTLY_RELEASE_DATE"):
   name += "-nightly"
   version += ".dev" + nightly_release_date
-
+  version_py.write_text(
+      re.sub(
+          version_regex, f'__version__ = "{version}"', version_py.read_text()
+      )
+  )
 
 setup(
     name=name,