fix DiagEmbedInferMeta when x_dims contains -1

Signed-off-by: ZelinMa557 <3388706467@qq.com>
PaddlePaddle · Apr 28, 2024 · 242b7b5 · 242b7b5
1 parent 71fd732
commit 242b7b5
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -840,7 +840,9 @@ void DiagEmbedInferMeta(
                         dim1,
                         dim2));
 
-  int new_dim_len = static_cast<int>(offset_ + x_dims[x_dims.size() - 1]);
+  int x_last_dim = x_dims[x_dims.size() - 1];
+  int new_dim_len =
+      (x_last_dim == -1) ? static_cast<int>(offset_ + x_last_dim) : -1;
   auto sizes = common::vectorize(x_dims);
   sizes.pop_back();
   sizes.insert(sizes.begin() + std::min(dim1_, dim2_), new_dim_len);