Fix roll_op by avoiding DivisionByZeroError, test=develop (#34499)

PaddlePaddle · Jul 30, 2021 · 5571c98 · 5571c98
1 parent ba19398
commit 5571c98
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/paddle/fluid/operators/roll_op.cu b/paddle/fluid/operators/roll_op.cu
@@ -80,9 +80,11 @@ class RollKernel<platform::CUDADeviceContext, T>
         int dim = dims[i] >= 0 ? dims[i] : dims[i] + input_dim.size();
         int64_t size = input_dim[dim];
 
-        shifts[i] = (shifts[i] % size + size) % size;
-        strides[i] = stride_dim[dim];
-        sizes[i] = size;
+        if (size != 0) {
+          shifts[i] = (shifts[i] % size + size) % size;
+          strides[i] = stride_dim[dim];
+          sizes[i] = size;
+        }
       }
     }
 
@@ -151,10 +153,11 @@ class RollGradKernel<platform::CUDADeviceContext, T>
       for (size_t i = 0; i < nums; i++) {
         int dim = dims[i] >= 0 ? dims[i] : dims[i] + input_dim.size();
         int64_t size = input_dim[dim];
-
-        shifts[i] = ((-shifts[i]) % size + size) % size;
-        strides[i] = stride_dim[dim];
-        sizes[i] = size;
+        if (size != 0) {
+          shifts[i] = ((-shifts[i]) % size + size) % size;
+          strides[i] = stride_dim[dim];
+          sizes[i] = size;
+        }
       }
     }
 

diff --git a/paddle/fluid/operators/roll_op.h b/paddle/fluid/operators/roll_op.h
@@ -30,6 +30,9 @@ inline void shift_along_dim(T* data, const DDim& input_dim, int64_t dim,
   if (dim < 0) {
     dim += input_dim.size();
   }
+  if (input_dim[dim] == 0) {
+    return;
+  }
   shift = shift % input_dim[dim];
   if (shift < 0) {
     shift += input_dim[dim];