refine unsqueeze_kernel (PaddlePaddle#58275)

wanghuancoder · jiahy0825 · commit fba8f03a2fc9 · 2023-10-26T04:19:09.000Z
diff --git a/paddle/phi/kernels/unsqueeze_kernel.cc b/paddle/phi/kernels/unsqueeze_kernel.cc
@@ -27,7 +27,7 @@ void UnsqueezeInferKernel(const Context& dev_ctx,
                           DenseTensor* out) {
   auto x_dims = x.dims();
   auto out_dims = out->dims();
-  if (axes.FromTensor()) {
+  if (axes.FromTensor() && out->dims()[0] == -1) {
     out_dims = funcs::GetUnsqueezeShape(axes.GetData(), x_dims);
   }
   out->Resize(out_dims);

Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,7 @@ void UnsqueezeInferKernel(const Context& dev_ctx,`
`27`	`27`	`DenseTensor* out) {`
`28`	`28`	`auto x_dims = x.dims();`
`29`	`29`	`auto out_dims = out->dims();`
`30`		`- if (axes.FromTensor()) {`
	`30`	`+ if (axes.FromTensor() && out->dims()[0] == -1) {`
`31`	`31`	`out_dims = funcs::GetUnsqueezeShape(axes.GetData(), x_dims);`
`32`	`32`	`}`
`33`	`33`	`out->Resize(out_dims);`