[Test] enable scaled-int dtypes for pruning test for inf cost

fastmachinelearning · Jan 30, 2025 · 5b050a7 · 5b050a7
1 parent ce063a9
commit 5b050a7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/transformation/test_pruning.py b/tests/transformation/test_pruning.py
@@ -86,7 +86,7 @@ def test_pruning_mnv1():
     model = download_model("MobileNetv1-w4a4", return_modelwrapper=True)
     # mark input as scaled 8-bit to get correct inference cost
     model.set_tensor_datatype(model.graph.input[0].name, DataType["SCALEDINT<8>"])
-    model = model.transform(InferDataTypes())
+    model = model.transform(InferDataTypes(allow_scaledint_dtypes=True))
     # do cleanup including folding quantized weights
     model = cleanup_model(model, False)
     inp, golden = get_golden_in_and_output("MobileNetv1-w4a4")