From 566c09bfa2409fb62357f6bdc9eac606291131ab Mon Sep 17 00:00:00 2001
From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com>
Date: Mon, 12 Oct 2026 11:56:22 -0700
Subject: [PATCH] add export_backend param

---
 python/src/coreai_models/export/compression.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/python/src/coreai_models/export/compression.py b/python/src/coreai_models/export/compression.py
index f29beac0..161795ae 100644
--- a/python/src/coreai_models/export/compression.py
+++ b/python/src/coreai_models/export/compression.py
@@ -326,6 +326,7 @@ def palettize_pytorch_model(
     model: nn.Module,
     example_inputs: tuple,
     palettization_config: "dict | KMeansPalettizerConfig",
+    export_backend: ExportBackend = ExportBackend.CoreAI,
 ) -> nn.Module:
     """
     Palettize a PyTorch model using post-training palettization with coreai-opt.
@@ -337,6 +338,7 @@ def palettize_pytorch_model(
         palettization_config: Either a configuration dictionary (matching the
             inner shape coreai-opt expects under `kmeans_palettization_config`)
             or a prebuilt KMeansPalettizerConfig instance.
+        export_backend: Backend for the finalized palettized model.
 
     Returns:
         Palettized model ready for export.
@@ -359,7 +361,7 @@ def palettize_pytorch_model(
     palettizer = KMeansPalettizer(model, config)
     prepared_model = palettizer.prepare(example_inputs=example_inputs, num_workers=32)
 
-    finalized_model = palettizer.finalize(prepared_model, backend=ExportBackend.CoreAI)
+    finalized_model = palettizer.finalize(prepared_model, backend=export_backend)
 
     logger.info("Palettization with coreai-opt complete")
     return finalized_model