support L2Norm backward for GPU. (apache#10874)

* support gpu. * Fix.
rahul003 · May 13, 2018 · 8b53a3d · 8b53a3d
1 parent ce57b58
commit 8b53a3d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/operator/tensor/broadcast_reduce_op_value.cu b/src/operator/tensor/broadcast_reduce_op_value.cu
@@ -104,5 +104,9 @@ NNVM_REGISTER_OP(norm)
 .set_attr<FCompute>("FCompute<gpu>", L2NormCompute<gpu>)
 .set_attr<FComputeEx>("FComputeEx<gpu>", L2NormComputeEx<gpu>);
 
+NNVM_REGISTER_OP(_backward_norm)
+.set_attr<FCompute>("FCompute<gpu>", L2NormGradCompute<gpu>)
+.set_attr<FComputeEx>("FComputeEx<gpu>", L2NormGradComputeEx<gpu>);
+
 }  // namespace op
 }  // namespace mxnet