From 2bf54bcaa2e299af56f8d4eeb8b94621cac14ac1 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Thu, 25 Mar 2021 14:53:52 -0700 Subject: [PATCH] Fix bugs in sparsify script (#7134) Fix type and check. --- tools/python/sparsify_initializers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/python/sparsify_initializers.py b/tools/python/sparsify_initializers.py index 61e9476dcf..c81d25c78a 100644 --- a/tools/python/sparsify_initializers.py +++ b/tools/python/sparsify_initializers.py @@ -16,7 +16,7 @@ from onnx import ModelProto, SparseTensorProto, TensorProto, numpy_helper logger = logging.getLogger(__name__) -real_types = set((np.float32, np.float64, np.double)) +real_types = set((int(TensorProto.FLOAT), int(TensorProto.DOUBLE))) def parse_arguments(): @@ -67,7 +67,7 @@ def convert_tensor_to_sparse(tensor, tolerance): # type: (TensorProto) -> Tuple else: for index in range(data_len): el = tensor_data[index] - if el == 0: + if el != 0: values.append(el) indicies.append(index) nnz_count += 1