From e073313d05a77806397d17bcaefef3ea633e3259 Mon Sep 17 00:00:00 2001
From: Yuanle Liu <yuanlehome@163.com>
Date: Fri, 23 Dec 2022 19:41:38 +0800
Subject: [PATCH] Fix arange gpu kernel (#49273)

---
 paddle/phi/kernels/gpu/arange_kernel.cu            | 9 +++++----
 python/paddle/fluid/tests/unittests/test_arange.py | 6 ++++++
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/paddle/phi/kernels/gpu/arange_kernel.cu b/paddle/phi/kernels/gpu/arange_kernel.cu
index 26f85446deb..4cd8eb3ceca 100644
--- a/paddle/phi/kernels/gpu/arange_kernel.cu
+++ b/paddle/phi/kernels/gpu/arange_kernel.cu
@@ -57,10 +57,11 @@ void ArangeKernel(const Context& dev_ctx,
   T* out_data = dev_ctx.template Alloc<T>(out);
 
   auto stream = dev_ctx.stream();
-  int block = std::min(size, static_cast<int64_t>(256));
-  PADDLE_ENFORCE_NE(
-      block, 0, errors::OutOfRange("The value of block cannot be 0."));
-  int grid = (size + block - 1) / block;
+  int64_t block = std::min(size, static_cast<int64_t>(256));
+  if (block == 0) {
+    return;
+  }
+  int64_t grid = (size + block - 1) / block;
   Range<T><<<grid, block, 0, stream>>>(start_value, step_value, size, out_data);
 }
 
diff --git a/python/paddle/fluid/tests/unittests/test_arange.py b/python/paddle/fluid/tests/unittests/test_arange.py
index f616acb0149..77c539172cd 100644
--- a/python/paddle/fluid/tests/unittests/test_arange.py
+++ b/python/paddle/fluid/tests/unittests/test_arange.py
@@ -70,6 +70,12 @@ class TestInt64ArangeOp(TestArangeOp):
         self.case = (-1, -10, -2)
 
 
+class TestZeroSizeArangeOp(TestArangeOp):
+    def init_config(self):
+        self.dtype = np.int32
+        self.case = (0, 0, 1)
+
+
 class TestArangeOpError(unittest.TestCase):
     def test_errors(self):
         with program_guard(Program(), Program()):
-- 
GitLab