From 454fab7448d4a23b0980abf94e2cd222145bad0e Mon Sep 17 00:00:00 2001 From: "Kwabena W. Agyeman" Date: Tue, 9 Jul 2024 22:19:15 -0700 Subject: [PATCH] modules/py_ml: Fix the size of the input bytearray. --- src/omv/modules/py_ml.c | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/omv/modules/py_ml.c b/src/omv/modules/py_ml.c index 91ab982e6..b27485289 100644 --- a/src/omv/modules/py_ml.c +++ b/src/omv/modules/py_ml.c @@ -37,6 +37,17 @@ static size_t py_ml_tuple_sum(mp_obj_tuple_t *o) { return size; } +static size_t pl_ml_dtype_size(char dtype) { + switch (dtype) { + case 'f': + return 4; + case 'h': + return 2; + default: + return 1; + } +} + static void py_ml_process_input(py_ml_model_obj_t *model, mp_obj_t arg) { mp_obj_list_t *input_list = MP_OBJ_TO_PTR(arg); @@ -49,7 +60,7 @@ static void py_ml_process_input(py_ml_model_obj_t *model, mp_obj_t arg) { if (mp_obj_is_callable(input_arg)) { // Input is a callable. Call the object and pass the tensor buffer and dtype. mp_obj_t fargs[3] = { - mp_obj_new_bytearray_by_ref(input_size, input_buffer), + mp_obj_new_bytearray_by_ref(input_size * pl_ml_dtype_size(model->input_dtype), input_buffer), MP_OBJ_FROM_PTR(input_shape), mp_obj_new_int(model->input_dtype) };