diff --git a/docker/build_image.py b/docker/build_image.py
index 231eb685..5f253eae 100644
--- a/docker/build_image.py
+++ b/docker/build_image.py
@@ -129,13 +129,8 @@ class BaseGPUImageBuilder(Builder):
         return content
 
     def image(self) -> str:
-        # return (
-        #     f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-{self.args.python_tag}-'
-        #     f'torch{self.args.torch_version}-tf{self.args.tf_version}-base')
-
-        # TODO: set the python tag to `py310` for base image temporarily
         return (
-            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-py310-'
+            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-{self.args.python_tag}-'
             f'torch{self.args.torch_version}-tf{self.args.tf_version}-base')
 
     def build(self) -> int: