From c4a68e48cc83f099340b1d7b375cc694b1f057d6 Mon Sep 17 00:00:00 2001 From: jfrery Date: Wed, 20 Nov 2024 16:47:54 +0100 Subject: [PATCH] chore: fix patch --- tests/torch/test_hybrid_converter.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/torch/test_hybrid_converter.py b/tests/torch/test_hybrid_converter.py index 110d5d5b2..67af03037 100644 --- a/tests/torch/test_hybrid_converter.py +++ b/tests/torch/test_hybrid_converter.py @@ -1,5 +1,6 @@ """Tests for the hybrid model converter.""" +import importlib import sys import tempfile from pathlib import Path @@ -64,14 +65,14 @@ def run_hybrid_llm_test( if has_pbs_reshape: has_pbs = True - # Propagate glwe_backend_installed state being tested to constants of affected modules - for affected_module in ( - concrete.ml.quantization.linear_op_glwe_backend, - concrete.ml.torch.hybrid_model, - ): - m.setattr(affected_module, "_HAS_GLWE_BACKEND", glwe_backend_installed) + # Patching for GLWE backend + if not glwe_backend_installed: + m.setitem(sys.modules, "concrete_ml_extensions", None) + + # Reload the affected modules to ensure the changes take effect + importlib.reload(concrete.ml.quantization.linear_op_glwe_backend) + importlib.reload(concrete.ml.torch.hybrid_model) - # Create a hybrid model hybrid_model = HybridFHEModel(model, module_names) is_compiled = False try: