chore: fix torch simulation

zama-ai · Apr 17, 2024 · 71b3f0f · 71b3f0f
1 parent f1172ad
commit 71b3f0f
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/use_case_examples/cifar/cifar_brevitas_finetuning/cifar_utils.py b/use_case_examples/cifar/cifar_brevitas_finetuning/cifar_utils.py
@@ -459,6 +459,8 @@ def fhe_compatibility(model: Callable, data: DataLoader) -> Callable:
         torch_inputset=data,
         show_mlir=False,
         output_onnx_file="test.onnx",
+        # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4391
+        rounding_threshold_bits=6,
     )
 
     return qmodel

diff --git a/use_case_examples/cifar/cifar_brevitas_training/Makefile b/use_case_examples/cifar/cifar_brevitas_training/Makefile
@@ -11,7 +11,5 @@ run_example: one two
 one:
 	@python evaluate_one_example_fhe.py
 
-# FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4389
-# too slow on cpu
-# two:
-# 	@python evaluate_torch_cml.py
+two:
+	@python evaluate_torch_cml.py
diff --git a/use_case_examples/cifar/cifar_brevitas_training/evaluate_torch_cml.py b/use_case_examples/cifar/cifar_brevitas_training/evaluate_torch_cml.py
@@ -139,7 +139,7 @@ def main(args):
         "--rounding_threshold_bits",
         nargs="+",
         type=int,
-        default=[None],
+        default=[6],
         help="Number of bits to target with rounding.",
     )
     parser.add_argument(