diff --git a/src/operator/contrib/nnz.cc b/src/operator/contrib/nnz.cc
index 956b22264667..a94f52fdaa5e 100644
--- a/src/operator/contrib/nnz.cc
+++ b/src/operator/contrib/nnz.cc
@@ -182,7 +182,8 @@ This operator only supports CSR matrix on CPU.
 .set_attr<nnvm::FInferType>("FInferType", NNZType)
 .set_attr<FInferStorageType>("FInferStorageType", NNZStorageType)
 .set_attr<FComputeEx>("FComputeEx<cpu>", NNZComputeEx<cpu>)
-.add_argument("data", "NDArray-or-Symbol", "Input");
+.add_argument("data", "NDArray-or-Symbol", "Input")
+.add_arguments(NNZParam::__FIELDS__());
 
 }  // namespace op
 }  // namespace mxnet