fix dtype compare bug

inisis · Nov 12, 2024 · 25f3eb6 · 25f3eb6
1 parent 2e05a2a
commit 25f3eb6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onnxslim/core/__init__.py b/onnxslim/core/__init__.py
@@ -171,7 +171,7 @@ def convert_data_format(model: onnx.ModelProto, dtype: str) -> onnx.ModelProto:
         for node in graph.nodes:
             if node.op == "Cast":
                 inp_dtype = [input.dtype for input in node.inputs][0]
-                if inp_dtype in {np.float16, np.float32}:
+                if inp_dtype in [np.float16, np.float32]:
                     delete_node(node)
 
         for tensor in graph.tensors().values():