aboutsummaryrefslogtreecommitdiff
path: root/nerv/matrix/cumatrix.c
diff options
context:
space:
mode:
Diffstat (limited to 'nerv/matrix/cumatrix.c')
-rw-r--r--nerv/matrix/cumatrix.c5
1 files changed, 3 insertions, 2 deletions
diff --git a/nerv/matrix/cumatrix.c b/nerv/matrix/cumatrix.c
index f6a4ed5..b8eef9c 100644
--- a/nerv/matrix/cumatrix.c
+++ b/nerv/matrix/cumatrix.c
@@ -9,7 +9,7 @@ const char *nerv_cuda_context_tname = "nerv.CuContext";
int nerv_cuda_context_lua_select_gpu(lua_State *L) {
Status status;
nerv_cuda_context_select_gpu(luaT_checkudata(L, 1, nerv_cuda_context_tname),
- luaL_checkinteger(L, 1), &status);
+ luaL_checkinteger(L, 2), &status);
NERV_LUA_CHECK_STATUS(L, status);
return 0;
}
@@ -26,7 +26,8 @@ int nerv_cuda_context_lua_clear_profile(lua_State *L) {
int nerv_cuda_context_lua_new(lua_State *L) {
Status status;
- CuContext *self = nerv_cuda_context_create(&status);
+ int dev = lua_gettop(L) > 0 ? luaL_checkinteger(L, 1) : -1;
+ CuContext *self = nerv_cuda_context_create(dev, &status);
NERV_LUA_CHECK_STATUS(L, status);
luaT_pushudata(L, self, nerv_cuda_context_tname);
return 1;