diff options
Diffstat (limited to 'compute/ARMComputeEx/arm_compute/runtime/CPP/functions/CPPOneHotEx.h')
-rw-r--r-- | compute/ARMComputeEx/arm_compute/runtime/CPP/functions/CPPOneHotEx.h | 10 |
1 files changed, 5 insertions, 5 deletions
diff --git a/compute/ARMComputeEx/arm_compute/runtime/CPP/functions/CPPOneHotEx.h b/compute/ARMComputeEx/arm_compute/runtime/CPP/functions/CPPOneHotEx.h index f132dfcc8..7930e4e20 100644 --- a/compute/ARMComputeEx/arm_compute/runtime/CPP/functions/CPPOneHotEx.h +++ b/compute/ARMComputeEx/arm_compute/runtime/CPP/functions/CPPOneHotEx.h @@ -55,14 +55,14 @@ public: /** Configure the one_hot function * * @param[in] indices A tensor for indices. Data types supported: S32 + * @param[in] depth A tensor for depth. Data types supported: S32 + * @param[in] on_value A tensor for on_value. Data types supported: F32 + * @param[in] off_value A tensor for off_value. Data types supported: F32 * @param[out] output A tensor for computed value of one hot operator - * @param[in] depth An int value for depth - * @param[in] on_value A float value for on_value - * @param[in] off_value A float value for off_value * @param[in] axis An int value for axis */ - void configure(const ITensor *indices, ITensor *output, const int depth, const float on_value, - const float off_value, const int axis); + void configure(const ITensor *indices, const ITensor *depth, const ITensor *on_value, + const ITensor *off_value, ITensor *output, const int axis); }; } #endif /* __ARM_COMPUTE_CPPONEHOT_EX_H__ */ |