diff --git a/paddle/fluid/framework/details/nan_inf_utils_detail.cu b/paddle/fluid/framework/details/nan_inf_utils_detail.cu
index 629ab737055a47c9493c686f982e0c05cf7441e2..4056fdc9cfe0e76fe88b0f37f1e9f91f564897b0 100644
--- a/paddle/fluid/framework/details/nan_inf_utils_detail.cu
+++ b/paddle/fluid/framework/details/nan_inf_utils_detail.cu
@@ -428,13 +428,12 @@ void TensorCheckerVisitor<phi::GPUContext>::apply(
   phi::DenseTensor block_num_nan_inf;
   block_num_nan_inf.Resize({static_cast<int64_t>(2 * numel_max_min)});
   int64_t* block_num_nan_ptr =
-      block_num_nan_inf.mutable_data<int64_t>(tensor.place());
+      dev_ctx->template Alloc<int64_t>(&block_num_nan_inf);
   int64_t* block_num_inf_ptr = block_num_nan_ptr + numel_max_min;
 
   phi::DenseTensor tensor_block_max_min;
   tensor_block_max_min.Resize({static_cast<int64_t>(3 * numel_max_min)});
-  MT* tensor_block_max_ptr =
-      tensor_block_max_min.mutable_data<MT>(tensor.place());
+  MT* tensor_block_max_ptr = dev_ctx->template Alloc<MT>(&tensor_block_max_min);
   MT* tensor_block_min_ptr = tensor_block_max_ptr + numel_max_min;
   MT* tensor_block_mean_ptr = tensor_block_max_ptr + 2 * numel_max_min;