From a69470135a717e02b7f5d53701f45258dfdd8ca0 Mon Sep 17 00:00:00 2001 From: Robert Sachunsky Date: Tue, 11 Oct 2022 07:21:17 +0000 Subject: [PATCH] load_model: also allow SavedModel directories --- sbb_binarize/sbb_binarize.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sbb_binarize/sbb_binarize.py b/sbb_binarize/sbb_binarize.py index 8960354..247d54b 100644 --- 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}")