improve test coverage

Signed-off-by: Gal Hubara Agam <[email protected]>

Author: galagam

tests/unit/onnx/autocast/test_nodeclassifier.py

@@ -369,3 +369,16 @@ def test_node_classifier_force_include(test_model):
     assert "add_node" in fp16_nodes
     assert "add_node" not in fp32_nodes
     assert not set(fp16_nodes).intersection(set(fp32_nodes))
+
+    # Set init_max low so both nodes would normally be excluded (kept in FP32)
+    # Force op type Mul to low precision, despite exceeding init_max
+    classifier3 = NodeClassifier(
+        model=test_model,
+        node_to_init_map=node_to_init_map,
+        init_max=1.0,
+        op_types_to_include=["Mul"],
+    )
+    fp16_nodes, fp32_nodes = classifier3.run()
+    assert "mul_node" in fp16_nodes
+    assert "add_node" in fp32_nodes
+    assert not set(fp16_nodes).intersection(set(fp32_nodes))