cpu_kernel_multiple_outputs Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/Loops.h lines 330–339

template <typename func_t>
void cpu_kernel_multiple_outputs(TensorIteratorBase& iter, func_t&& op, int64_t grain_size = at::internal::GRAIN_SIZE) {
  using traits = function_traits<func_t>;
  TORCH_INTERNAL_ASSERT(iter.ninputs() == traits::arity);

  iter.for_each([&](char** data, const int64_t* strides, int64_t n) {
    multiple_outputs_loop(data, strides, 0, n, op);
  }, grain_size);
  iter.cast_outputs();
}