diff --git a/paddle/fluid/prim/api/composite_backward/composite_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
index 35059ab2ddfbd96ae1635c5427b88bfbe49da989..58617f9b36d74a179cc9072d9ed2dbbe2560f33f 100644
--- a/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
+++ b/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
@@ -1299,7 +1299,7 @@ void batch_norm_grad(const Tensor& x,
 
   std::vector<int> nchw_to_nhwc_dim = {0, 2, 3, 1};
   std::vector<int> nhwc_to_nchw_dim = {0, 3, 1, 2};
-  auto reduce_axis = IntArray(std::vector<int>{0, 1, 2});
+  auto reduce_axis = IntArray(std::vector<int64_t>{0, 1, 2});
   auto dtype = x_data.dtype();
 
   switch (data_layout_) {