IS_INPUT Class — pytorch Architecture

Source Code

aten/src/ATen/cuda/jiterator_impl.h lines 77–92

template<bool IS_INPUT, int N>
static std::unique_ptr<OffsetCalculator<N>> make_unique_offset_calculator(
          const TensorIteratorBase& iter) {
  // array size can not be 0, this happens when N == 0
  constexpr int array_size = std::max<int>(N, 1);
  TORCH_INTERNAL_ASSERT(N == (IS_INPUT ? iter.ninputs() : iter.noutputs()));

  std::array<const int64_t*, array_size> strides;
  int64_t element_sizes[array_size];
  for (int i = 0; i < N; i++) {
    int index = IS_INPUT ? i + iter.noutputs() : i;
    strides[i] = iter.strides(index).data();
    element_sizes[i] = iter.element_size(index);
  }
  return std::make_unique<OffsetCalculator<N>>(iter.ndim(), iter.shape().data(), strides.data(), element_sizes);
}