xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/mobile/model_tracer/TensorUtils.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <ATen/core/ivalue.h>
4 
5 namespace torch::jit::mobile {
6 /**
7  * Recursively scan the IValue object, traversing lists, tuples, dicts, and stop
8  * and call the user provided callback function 'func' when a Tensor is found.
9  */
10 void for_each_tensor_in_ivalue(
11     const ::c10::IValue& iv,
12     std::function<void(const ::at::Tensor&)> const& func);
13 } // namespace torch::jit::mobile
14