Add gpu compile

eric-haibin-lin · Aug 3, 2017 · 9da68f3 · 9da68f3
1 parent 5c4e9c1
commit 9da68f3
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/src/c_api/c_api_ndarray.cc b/src/c_api/c_api_ndarray.cc
@@ -393,7 +393,7 @@ void PushOperator(const OpStatePtr& state,
 #if MXNET_USE_CUDA
           CastNonDefaultStorage<gpu>(temp_in_src, temp_in_dst, opctx);
           fcompute(state, opctx, input_blobs, req, output_blobs);
-          CastNonDefaultStorage<gpu>(temp_our_dst, temp_out_src, opctx);
+          CastNonDefaultStorage<gpu>(temp_out_dst, temp_out_src, opctx);
 #else
           LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
 #endif

diff --git a/src/operator/tensor/sparse_retain.cu b/src/operator/tensor/sparse_retain.cu
@@ -0,0 +1,18 @@
+/*!
+ * Copyright (c) 2017 by Contributors
+ * \file sparse_retain.cu
+ * \brief
+*/
+
+#include "./sparse_retain-inl.h"
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(sparse_retain)
+.set_attr<FComputeEx>("FComputeEx<gpu>", SparseRetainOpForwardEx<gpu>);
+
+NNVM_REGISTER_OP(_backward_sparse_retain)
+.set_attr<FComputeEx>("FComputeEx<gpu>", SparseRetainOpBackwardEx<gpu>);
+
+}  // namespace op
+}  // namespace mxnet