diff --git a/setup.py b/setup.py
index ce29c7462..39f8689fb 100644
--- a/setup.py
+++ b/setup.py
@@ -14,10 +14,9 @@
 
 
 def get_version():
-    # get version string from version.py
-    # TODO: ideally the version.py should be generated when setup is run
-    version_file = os.path.join(os.path.dirname(__file__), "torchx/version.py")
-    version_regex = r"__version__ = ['\"]([^'\"]*)['\"]"
+    # get version string from _version.py
+    version_file = os.path.join(os.path.dirname(__file__), "torchx/_version.py")
+    version_regex = r"BASE_VERSION = ['\"]([^'\"]*)['\"]"
     with open(version_file, "r") as f:
         version = re.search(version_regex, f.read(), re.M).group(1)
         return version
diff --git a/torchx/_version.py b/torchx/_version.py
new file mode 100644
index 000000000..09810e2c1
--- /dev/null
+++ b/torchx/_version.py
@@ -0,0 +1,8 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+BASE_VERSION = "0.8.0dev0"
diff --git a/torchx/version.py b/torchx/version.py
index e3a7632f7..746df34a5 100644
--- a/torchx/version.py
+++ b/torchx/version.py
@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -7,6 +6,7 @@
 
 # pyre-strict
 
+from torchx._version import BASE_VERSION
 from torchx.util.entrypoints import load
 
 # Follows PEP-0440 version scheme guidelines
@@ -18,7 +18,7 @@
 # 0.1.0bN  # Beta release
 # 0.1.0rcN  # Release Candidate
 # 0.1.0  # Final release
-__version__ = "0.8.0dev0"
+__version__: str = BASE_VERSION
 
 
 # Use the github container registry images corresponding to the current package