group dim

Xilinx · Dec 19, 2024 · 4dad221 · 4dad221
1 parent 2ac05a1
commit 4dad221
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/brevitas/export/inference/handler.py b/src/brevitas/export/inference/handler.py
@@ -287,6 +287,7 @@ def prepare_for_export(self, module: nn.Module):
         if module.is_quant_enabled:
             self.input_view = module.input_view_impl
             self.flattened_view = module.apply_input_view
+            self.group_dim = module.group_dim
             if module._cached_weight is not None and not module.cache_inference_quant_weight_metadata_only:
                 self.cached_weight = module._cached_weight.quant_tensor.value_
             else: