Fix bugs in sparsify script (#7134)

Fix type and check.
This commit is contained in:
Dmitri Smirnov 2021-03-25 14:53:52 -07:00 committed by GitHub
parent cc0e7bee76
commit 2bf54bcaa2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -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