chore: remove tf warnings at inference

zama-ai · Sep 29, 2023 · 920677b · 920677b
1 parent b3f892d
commit 920677b
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/use_case_examples/hybrid_model/infer_hybrid_llm_generate.py b/use_case_examples/hybrid_model/infer_hybrid_llm_generate.py
@@ -1,5 +1,6 @@
 """Showcase for the hybrid model converter."""
 import json
+import os
 import time
 from pathlib import Path
 
@@ -9,6 +10,12 @@
 
 from concrete.ml.torch.hybrid_model import HybridFHEMode, HybridFHEModel
 
+os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+import tensorflow
+
+tensorflow.compat.v1.logging.set_verbosity(tensorflow.compat.v1.logging.ERROR)
+
 if __name__ == "__main__":
     # Load configuration dumped by compilation
     with open("configuration.json", "r") as file: