Fix

Author: Giuseppe5

src/brevitas/graph/equalize.py

@@ -922,6 +922,8 @@ def forward_stats_hook(self, module, *args, name, batch_dim=0, use_inp=True, **k
         if hasattr(x, 'names') and 'N' in x.names:
             batch_dim = x.names.index('N')
 
+        self.batch_dim_act_map[name] = batch_dim
+
         input_scales = self.scale_fn(x, dim=batch_dim)
         if name not in self.float_act_map:
             self.float_act_map[name] = input_scales

src/brevitas_examples/llm/llm_quant/equalize.py

@@ -31,7 +31,6 @@ def activation_equalization_iter(curr_layer, inps, outs, cached_values, alpha):
 def trace_and_standardize(model, ref_kwargs):
     graph_model = value_trace(model, value_args=ref_kwargs)
     graph_model = TorchFunctionalToModule().apply(graph_model)
-    graph_model = DuplicateSharedStatelessModule().apply(graph_model)
     return graph_model