Fix (brevitas_examples/sdxl): correct import for inference_mode (#1185)

Xilinx · Feb 13, 2025 · 39ccbde · 39ccbde
1 parent db8af71
commit 39ccbde
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/brevitas_examples/stable_diffusion/main.py b/src/brevitas_examples/stable_diffusion/main.py
@@ -27,6 +27,7 @@
 from tqdm import tqdm
 
 from brevitas.core.stats.stats_op import NegativeMinOrZero
+from brevitas.export.inference import quant_inference_mode
 from brevitas.export.onnx.standard.qcdq.manager import StdQCDQONNXManager
 from brevitas.graph.base import ModuleToModuleByClass
 from brevitas.graph.calibrate import bias_correction_mode
@@ -43,7 +44,6 @@
 from brevitas_examples.common.parse_utils import add_bool_arg
 from brevitas_examples.common.parse_utils import quant_format_validator
 from brevitas_examples.llm.llm_quant.export import BlockQuantProxyLevelManager
-from brevitas_examples.llm.main import quant_inference_mode
 from brevitas_examples.stable_diffusion.mlperf_evaluation.accuracy import compute_mlperf_fid
 from brevitas_examples.stable_diffusion.sd_quant.constants import SD_2_1_EMBEDDINGS_SHAPE
 from brevitas_examples.stable_diffusion.sd_quant.constants import SD_XL_EMBEDDINGS_SHAPE