construct only necessary elements in OffsetCalculator

pytorch · Mar 31, 2021 · 646510f · 646510f
1 parent 4865195
commit 646510f
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/aten/src/ATen/cuda/detail/OffsetCalculator.cuh b/aten/src/ATen/cuda/detail/OffsetCalculator.cuh
@@ -28,15 +28,11 @@ struct OffsetCalculator {
   // the strides will be in # of elements.
   OffsetCalculator(int dims, const int64_t* sizes, const int64_t* const* strides, const int64_t* element_sizes=nullptr) : dims(dims) {
     TORCH_CHECK(dims <= MAX_DIMS, "tensor has too many (>", MAX_DIMS, ") dims");
-    for (int i = 0; i < MAX_DIMS; ++i) {
-      if (i < dims) {
-        sizes_[i] = IntDivider<index_t>(sizes[i]);
-      } else {
-        sizes_[i] = IntDivider<index_t>(1);
-      }
+    for (int i=0; i < dims; i++){
+      sizes_[i] = IntDivider<index_t>(sizes[i]);
       for (int arg = 0; arg < NARGS; arg++) {
         int64_t element_size = (element_sizes == nullptr ? 1LL : element_sizes[arg]);
-        strides_[i][arg] =  i < dims ? strides[arg][i] / element_size : 0;
+        strides_[i][arg] = strides[arg][i] / element_size;
       }
     }
   }