add Register depth_to_space in operator.cc

mace/core/operator.cc

@@ -73,6 +73,7 @@ extern void Register_BufferToImage(OperatorRegistry *op_registry);
 extern void Register_ChannelShuffle(OperatorRegistry *op_registry);
 extern void Register_Concat(OperatorRegistry *op_registry);
 extern void Register_Conv2D(OperatorRegistry *op_registry);
+extern void Register_DepthToSpace(OperatorRegistry *op_registry);
 extern void Register_DepthwiseConv2d(OperatorRegistry *op_registry);
 extern void Register_FoldedBatchNorm(OperatorRegistry *op_registry);
 extern void Register_FusedConv2D(OperatorRegistry *op_registry);
@@ -103,6 +104,7 @@ OperatorRegistry::OperatorRegistry() {
   ops::Register_Concat(this);
   ops::Register_Conv2D(this);
   ops::Register_DepthwiseConv2d(this);
+  ops::Register_DepthToSpace(this);
   ops::Register_FoldedBatchNorm(this);
   ops::Register_FusedConv2D(this);
   ops::Register_GlobalAvgPooling(this);

mace/ops/depth_to_space.cc

@@ -8,11 +8,18 @@ namespace mace {
 namespace ops {
 
 void Register_DepthToSpace(OperatorRegistry *op_registry) {
+  REGISTER_OPERATOR(op_registry, OpKeyBuilder("DepthToSpace")
+                                     .Device(DeviceType::CPU)
+                                     .TypeConstraint<float>("T")
+                                     .Build(),
+                    DepthToSpaceOp<DeviceType::CPU, float>);
+                    
   REGISTER_OPERATOR(op_registry, OpKeyBuilder("DepthToSpace")
                                      .Device(DeviceType::OPENCL)
                                      .TypeConstraint<float>("T")
                                      .Build(),
                     DepthToSpaceOp<DeviceType::OPENCL, float>);
+                    
   REGISTER_OPERATOR(op_registry, OpKeyBuilder("DepthToSpace")
                                      .Device(DeviceType::OPENCL)
                                      .TypeConstraint<half>("T")

mace/ops/depth_to_space.h

@@ -20,7 +20,7 @@ class DepthToSpaceOp : public Operator<D, T> {
   DepthToSpaceOp(const OperatorDef &op_def, Workspace *ws)
       : Operator<D, T>(op_def, ws),
         functor_(OperatorBase::GetRepeatedArgument<int>("crops", {0, 0, 0, 0}),
-                 OperatorBase::GetRepeatedArgument<int>("block_shape", {1, 1}),
+                 OperatorBase::GetSingleArgument<int>("block_size", 1),
                  true) {}
 
   bool Run(StatsFuture *future) override {