memsize Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/SumKernel.cpp lines 197–209

template <typename scalar_t>
struct NanSumLoadPolicy<Vectorized<scalar_t>> {
  using vec_t = Vectorized<scalar_t>;

  static constexpr int64_t memsize() {
    return LoadPolicy<vec_t>::memsize();
  }

  static vec_t load(const char * C10_RESTRICT data, int64_t stride, int64_t index) {
    auto val = LoadPolicy<vec_t>::load(data, stride, index);
    return vec_t::blendv(val, vec_t(0), val.isnan());
  }
};