diff --git a/tensorflow/python/distribute/saved_model_test_base.py b/tensorflow/python/distribute/saved_model_test_base.py
index 245e258ffdb..31b84b13b88 100644
--- a/tensorflow/python/distribute/saved_model_test_base.py
+++ b/tensorflow/python/distribute/saved_model_test_base.py
@@ -19,7 +19,6 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import tempfile
 from absl.testing import parameterized
 import numpy as np
 
@@ -155,8 +154,7 @@ class TestSavedModelBase(test.TestCase, parameterized.TestCase):
                                                  distribution, run_distributed):
     """Save a model without DS, and restore it with DS."""
 
-    saved_dir = os.path.join(tempfile.mkdtemp(dir=self.get_temp_dir()),
-                             self._root_dir, 'test_save_no_dist_restore_dist')
+    saved_dir = os.path.join(self.get_temp_dir(), '0')
 
     model, output_name = model_and_input.get_model(
         run_distributed=run_distributed)
@@ -183,8 +181,7 @@ class TestSavedModelBase(test.TestCase, parameterized.TestCase):
                                                  run_distributed):
     """Save a model with DS, and restore it without DS."""
 
-    saved_dir = os.path.join(tempfile.mkdtemp(dir=self.get_temp_dir()),
-                             self._root_dir, 'test_save_no_dist_restore_dist')
+    saved_dir = os.path.join(self.get_temp_dir(), '1')
 
     with distribution.scope():
       model, output_name = model_and_input.get_model(
@@ -216,8 +213,7 @@ class TestSavedModelBase(test.TestCase, parameterized.TestCase):
                                               save_in_scope, run_distributed):
     """Save a model with DS, and restore it with potentially different DS."""
 
-    saved_dir = os.path.join(tempfile.mkdtemp(dir=self.get_temp_dir()),
-                             self._root_dir, 'test_save_dist_restore_dist')
+    saved_dir = os.path.join(self.get_temp_dir(), '2')
 
     with distribution_for_saving.scope():
       model, output_name = model_and_input.get_model(