load_model: also allow SavedModel directories

2 years ago · a69470135a
parent 88d64de67b
commit a69470135a
1 changed files with 2 additions and 0 deletions
--- a/sbb_binarize/sbb_binarize.py
+++ b/sbb_binarize/sbb_binarize.py
@ -34,6 +34,8 @@ class SbbBinarizer:
        self.start_new_session()
        self.model_files = glob('%s/*.h5' % self.model_dir)
        if not self.model_files:
            self.model_files = glob('%s/*/' % self.model_dir)
        if not self.model_files:
            raise ValueError(f"No models found in {self.model_dir}")