fix lp search test

sony · Jan 7, 2025 · dc0755e · dc0755e
1 parent 07109f5
commit dc0755e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/keras_tests/non_parallel_tests/test_lp_search_bitwidth.py b/tests/keras_tests/non_parallel_tests/test_lp_search_bitwidth.py
@@ -63,7 +63,7 @@ def __init__(self, layer_to_ru_mapping):
         self.compute_metric_fn = lambda x, y=None, z=None: {0: 2, 1: 1, 2: 0}[x[0]]
         self.min_ru = {RUTarget.WEIGHTS: [[1], [1], [1]],
                        RUTarget.ACTIVATION: [[1], [1], [1]],
-                       RUTarget.TOTAL: [[2], [2], [2]],
+                       RUTarget.TOTAL: [[1, 1], [1, 1], [1, 1]],
                        RUTarget.BOPS: [[1], [1], [1]]}  # minimal resource utilization in the tests layer_to_ru_mapping
 
         self.max_ru_config = [0]
@@ -77,8 +77,8 @@ def compute_resource_utilization_matrix(self, target):
         elif target == RUTarget.ACTIVATION:
             ru_matrix = [np.flip(np.array([ru.activation_memory - 1 for _, ru in self.layer_to_ru_mapping[0].items()]))]
         elif target == RUTarget.TOTAL:
-            ru_matrix = [np.flip(np.array([ru.weights_memory - 1 for _, ru in self.layer_to_ru_mapping[0].items()])),
-                         np.flip(np.array([ru.activation_memory - 1 for _, ru in self.layer_to_ru_mapping[0].items()]))]
+            ru_matrix = [[np.flip(np.array([ru.weights_memory - 1 for _, ru in self.layer_to_ru_mapping[0].items()])),
+                         np.flip(np.array([ru.activation_memory - 1 for _, ru in self.layer_to_ru_mapping[0].items()]))]]
         elif target == RUTarget.BOPS:
             ru_matrix = [np.flip(np.array([ru.bops - 1 for _, ru in self.layer_to_ru_mapping[0].items()]))]
         else: