changed thres_mask function of matrix to a more standard api

1 files changed, 5 insertions, 4 deletions

diff --git a/nerv/matrix/generic/cumatrix.c b/nerv/matrix/generic/cumatrix.c
index d1f763b..fb36033 100644
--- a/nerv/matrix/generic/cumatrix.c
+++ b/nerv/matrix/generic/cumatrix.c
@@ -65,10 +65,11 @@ static int nerv_matrix_(lua_sigmoid_grad)(lua_State *L) {
 static int nerv_matrix_(lua_thres_mask)(lua_State *L) {
     Status status;
     Matrix *a = luaT_checkudata(L, 1, nerv_matrix_(tname));
-    MATRIX_ELEM thres = luaL_checknumber(L, 2);
-    MATRIX_ELEM low = luaL_checknumber(L, 3);
-    MATRIX_ELEM high = luaL_checknumber(L, 4);
-    nerv_matrix_(thres_mask)(a, thres, low, high, &status);
+    Matrix *b = luaT_checkudata(L, 2, nerv_matrix_(tname));
+    MATRIX_ELEM thres = luaL_checknumber(L, 3);
+    MATRIX_ELEM low = luaL_checknumber(L, 4);
+    MATRIX_ELEM high = luaL_checknumber(L, 5);
+    nerv_matrix_(thres_mask)(a, b, thres, low, high, &status);
     NERV_LUA_CHECK_STATUS(L, status);
     return 0;
 }