Fix hard coded checkpoint prefix for tensorflow deployer. (#464)

ucbrise · Apr 13, 2018 · 4e775e3 · 4e775e3
1 parent 0d89cb2
commit 4e775e3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/clipper_admin/clipper_admin/deployers/tensorflow.py b/clipper_admin/clipper_admin/deployers/tensorflow.py
@@ -218,9 +218,9 @@ def predict(sess, inputs):
                     raise e
             else:
                 # Check if all the 4 files are present
-                if glob.glob('data/*.meta') and glob.glob(
-                        'data/*.index') and glob.glob(
-                            "data/checkpoint") and glob.glob("data/*.data*"):
+                suffixes = [os.path.join(tf_sess_or_saved_model_path, suffix) \
+                           for suffix in ("*.meta", "*.index", "checkpoint", "*.data*")]
+                if sum([1 for suffix in suffixes if glob.glob(suffix)]) == 4:
                     try:
                         shutil.copytree(tf_sess_or_saved_model_path,
                                         os.path.join(serialization_dir,