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) };