Fix for act eq

Xilinx · Jul 6, 2023 · a454afc · a454afc
1 parent 408439e
commit a454afc
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/brevitas_examples/imagenet_classification/ptq/benchmark/ptq_benchmark_torchvision.py b/src/brevitas_examples/imagenet_classification/ptq/benchmark/ptq_benchmark_torchvision.py
@@ -55,7 +55,7 @@
     'bias_corr': [True],  # Bias Correction
     'graph_eq_iterations': [0, 20],  # Graph Equalization
     'graph_eq_merge_bias': [False, True],  # Merge bias for Graph Equalization
-    'act_eq': ['fx', 'layerwise', None],  # Perform Activation Equalization (Smoothquant)
+    'act_equalization': ['fx', 'layerwise', None],  # Perform Activation Equalization (Smoothquant)
     'learned_round': [False, True],  # Enable/Disable Learned Round
     'gptq': [False, True],  # Enable/Disable GPTQ
     'gptq_act_order': [False, True],  # Use act_order euristics for GPTQ
@@ -73,7 +73,7 @@
     'bias_corr': [True],  # Bias Correction
     'graph_eq_iterations': [20],  # Graph Equalization
     'graph_eq_merge_bias': [True],  # Merge bias for Graph Equalization
-    'act_eq': ['fx'],  # Perform Activation Equalization (Smoothquant)
+    'act_equalization': ['fx'],  # Perform Activation Equalization (Smoothquant)
     'learned_round': [False],  # Enable/Disable Learned Round
     'gptq': [True],  # Enable/Disable GPTQ
     'gptq_act_order': [False],  # Use act_order euristics for GPTQ
@@ -204,7 +204,8 @@ def ptq_torchvision_models(df, args):
 
     if config_namespace.act_equalization is not None:
         print("Applying activation equalization:")
-        apply_act_equalization(model, calib_loader, layerwise=args.act_equalization == 'layerwise')
+        apply_act_equalization(
+            model, calib_loader, layerwise=config_namespace.act_equalization == 'layerwise')
 
     # Define the quantized model
     quant_model = quantize_model(