Minor fix in ILP missing import and sensitivity test

sony · Jan 15, 2025 · 35fca8f · 35fca8f
1 parent 259ff8e
commit 35fca8f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py b/model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py
@@ -16,7 +16,7 @@
 import numpy as np
 from pulp import *
 from tqdm import tqdm
-from typing import Dict, Tuple, Set, Any
+from typing import Dict, Tuple, Any, Optional
 
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization, RUTarget

diff --git a/tests/keras_tests/function_tests/test_sensitivity_metric_interest_points.py b/tests/keras_tests/function_tests/test_sensitivity_metric_interest_points.py
@@ -35,7 +35,7 @@
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import get_mp_interest_points
-from model_compression_toolkit.core import DEFAULTCONFIG
+from model_compression_toolkit.core import DEFAULTCONFIG, CustomOpsetLayers
 from model_compression_toolkit.core.common.quantization.set_node_quantization_config import \
     set_quantization_configuration_to_graph
 from model_compression_toolkit.core.common.similarity_analyzer import compute_mse, compute_kl_divergence
@@ -68,7 +68,7 @@ def dummy_representative_dataset():
                                                                       c.activation_n_bits) for c in mixed_precision_cfg_list],
                                         name="sem_test")
 
-    fqc = AttachTpcToKeras().attach(tpc, custom_opset2layer={"Input": ([InputLayer],)})
+    fqc = AttachTpcToKeras().attach(tpc, custom_opset2layer={"Input": CustomOpsetLayers([InputLayer])})
 
     graph.set_fqc(fqc)
     graph = set_quantization_configuration_to_graph(graph=graph,