add FListInputNames attribute to softmax_cross_entropy

apache · Sep 28, 2018 · 9f27699 · 9f27699
1 parent 9ce5792
commit 9f27699
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/operator/loss_binary_op.cc b/src/operator/loss_binary_op.cc
@@ -67,6 +67,10 @@ Example::
   })
 .set_attr<FCompute>("FCompute<cpu>", SoftmaxCrossEntropyForward<cpu>)
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_softmax_cross_entropy"})
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+  [](const NodeAttrs& attrs) {
+    return std::vector<std::string>{"data", "label"};
+})
 .add_argument("data", "NDArray-or-Symbol", "Input data")
 .add_argument("label", "NDArray-or-Symbol", "Input label");