| /*! | |
| * Copyright (c) 2016 by Contributors | |
| * \file init_op.cu | |
| * \brief GPU Implementation of init op | |
| */ | |
| #include "./init_op.h" | |
| namespace mxnet { | |
| namespace op { | |
| NNVM_REGISTER_OP(_zeros) | |
| .set_attr<FCompute>("FCompute<gpu>", FillCompute<gpu, 0>); | |
| NNVM_REGISTER_OP(_ones) | |
| .set_attr<FCompute>("FCompute<gpu>", FillCompute<gpu, 1>); | |
| NNVM_REGISTER_OP(_arange) | |
| .set_attr<FCompute>("FCompute<gpu>", RangeCompute<gpu>); | |
| } // namespace op | |
| } // namespace mxnet |