MLIA-469 Support batch size > 1 for optimizations

- Add a PruningPolicy to skip layers that are not
  supported by the Keras pruning API
- Make dataset generation more generic to support
  use-cases beyond classification

Signed-off-by: Annie Tallund <annie.tallund@arm.com>
Change-Id: I198dae2b53860f449f2fdbc71575babceed1ffcf

1 files changed, 1 insertions, 7 deletions

diff --git a/tests/test_nn_tensorflow_utils.py b/tests/test_nn_tensorflow_utils.py
index 5131171..14b06c4 100644
--- a/tests/test_nn_tensorflow_utils.py
+++ b/tests/test_nn_tensorflow_utils.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
 # SPDX-License-Identifier: Apache-2.0
 """Test for module utils/test_utils."""
 from pathlib import Path
@@ -31,12 +31,6 @@ def test_generate_representative_dataset() -> None:
         assert isinstance(ndarray, np.ndarray)
 
 
-def test_generate_representative_dataset_wrong_shape() -> None:
-    """Test that only shape with batch size=1 is supported."""
-    with pytest.raises(Exception, match="Only the input batch_size=1 is supported!"):
-        representative_dataset([2, 3, 3], 5)
-
-
 def test_convert_saved_model_to_tflite(test_tf_model: Path) -> None:
     """Test converting SavedModel to TensorFlow Lite."""
     result = convert_to_tflite(test_tf_model.as_posix())