From 8013914bdfefd1ed9375b0c097537fb80d38893d Mon Sep 17 00:00:00 2001
From: Ofir Gordon <Ofir.Gordon@altair-semi.com>
Date: Mon, 11 Mar 2024 13:14:29 +0200
Subject: [PATCH] fix custom layer test
---
.../custom_layers_tests/test_sony_ssd_postprocess_layer.py | 6 +++---
1 file changed, 3 insertions(+), 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
index 6e7f8d68d..6fe8181fc 100644
--- 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')