chore: print stats

zama-ai · Jun 5, 2024 · fa2a20f · fa2a20f
1 parent c35c57e
commit fa2a20f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/use_case_examples/cifar/cifar_brevitas_training/evaluate_one_example_fhe.py b/use_case_examples/cifar/cifar_brevitas_training/evaluate_one_example_fhe.py
@@ -4,6 +4,7 @@
 from functools import partial
 from importlib.metadata import version
 from pathlib import Path
+from pprint import pprint
 
 import numpy as np
 import torch
@@ -114,6 +115,7 @@ def wrapper(*args, **kwargs):
     p_error=P_ERROR,
 )
 assert isinstance(quantized_numpy_module, QuantizedModule)
+assert quantized_numpy_module.fhe_circuit is not None
 
 print(tlu_optimizer.statistics)
 
@@ -130,6 +132,8 @@ def wrapper(*args, **kwargs):
 open("cifar10.graph", "w").write(str(quantized_numpy_module.fhe_circuit))
 open("cifar10.mlir", "w").write(quantized_numpy_module.fhe_circuit.mlir)
 
+pprint(quantized_numpy_module.fhe_circuit.statistics)
+
 # import sys
 #
 # if sys.platform == "darwin":