fix custom layer test

sony · Mar 11, 2024 · 8013914 · 8013914
1 parent b4086fc
commit 8013914
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/keras_tests/custom_layers_tests/test_sony_ssd_postprocess_layer.py b/tests/keras_tests/custom_layers_tests/test_sony_ssd_postprocess_layer.py
@@ -49,11 +49,11 @@ def test_custom_layer(self):
 
         core_config = mct.core.CoreConfig(
             mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(
-                use_hessian_based_scores=False))
+                use_hessian_based_scores=False,
+                target_kpi=mct.core.KPI(weights_memory=6000)))
         q_model, _ = mct.ptq.keras_post_training_quantization(model,
                                                               get_rep_dataset(2, (1, 8, 8, 3)),
-                                                              core_config=core_config,
-                                                              target_kpi=mct.core.KPI(weights_memory=6000))
+                                                              core_config=core_config)
 
         # verify the custom layer is in the quantized model
         self.assertTrue(isinstance(q_model.layers[-1], SSDPostProcess), 'Custom layer should be in the quantized model')