Fix stack backward

oneflow/core/autograd/gradient_funcs/stack.cpp

@@ -65,14 +65,10 @@ Maybe<void> Stack::Apply(const StackCaptureState* ctx, const TensorTuple& out_gr
   in_grads->resize(ctx->input_num);
   TensorTuple like(ctx->input_num);
   for (int i = 0; i < ctx->input_num; ++i) { like[i] = ctx->SavedTensors().at(i); }
-  if (ctx->input_num == 1) {
-    in_grads->at(0) = out_grads.at(0);
-  } else {
-    const auto& results = JUST(functional::StackGrad(out_grads.at(0), like, ctx->axis));
-    CHECK_EQ_OR_RETURN(results->size(), ctx->input_num);
-    for (int i = 0; i < ctx->input_num; ++i) {
-      if (ctx->requires_grad.at(i)) { in_grads->at(i) = results->at(i); }
-    }
+  const auto& results = JUST(functional::StackGrad(out_grads.at(0), like, ctx->axis));
+  CHECK_EQ_OR_RETURN(results->size(), ctx->input_num);
+  for (int i = 0; i < ctx->input_num; ++i) {
+    if (ctx->requires_grad.at(i)) { in_grads->at(i) = results->at(i); }
   }
   return Maybe<void>::Ok();
 }

oneflow/core/functional/impl/array_functor.cpp

@@ -524,6 +524,7 @@ class StackFunctor {
     CHECK_OR_RETURN(stack_dim >= 0 && stack_dim <= ndims)
         << "Index Error: Dimension out of range (expected in range of [" << -ndims - 1 << ", "
         << ndims << "], but got " << stack_dim;
+    if (ninput == 1) { return ExpandDims(inputs[0], dim); }
     const std::shared_ptr<const Shape>& first_in_shape = inputs[0]->shape();
     for (const auto& input : inputs) {
       for (int i = 0; i < ndims; ++i) {