expose retain and free for CudaStorage/CudaTensor, as already done for cpu storages/tensors

2 files changed, 18 insertions, 0 deletions

diff --git a/torch/generic/Storage.c b/torch/generic/Storage.c
index 62e887b..27d7b06 100644
--- a/torch/generic/Storage.c
+++ b/torch/generic/Storage.c
@@ -63,6 +63,13 @@ static int torch_Storage_(new)(lua_State *L)
   return 1;
 }
 
+static int torch_Storage_(retain)(lua_State *L)
+{
+  THStorage *storage = luaT_checkudata(L, 1, torch_Storage);
+  THStorage_(retain)(cutorch_getstate(L), storage);
+  return 0;
+}
+
 static int torch_Storage_(free)(lua_State *L)
 {
   THStorage *storage = luaT_checkudata(L, 1, torch_Storage);
@@ -220,6 +227,8 @@ static int torch_Storage_(read)(lua_State *L)
 }
 
 static const struct luaL_Reg torch_Storage_(_) [] = {
+  {"retain", torch_Storage_(retain)},
+  {"free", torch_Storage_(free)},
   {"size", torch_Storage_(__len__)},
   {"__len__", torch_Storage_(__len__)},
   {"__newindex__", torch_Storage_(__newindex__)},
diff --git a/torch/generic/Tensor.c b/torch/generic/Tensor.c
index 136dae1..f56c5b7 100644
--- a/torch/generic/Tensor.c
+++ b/torch/generic/Tensor.c
@@ -928,6 +928,13 @@ static int torch_Tensor_(__index__)(lua_State *L)
   }
 }
 
+static int torch_Tensor_(retain)(lua_State *L)
+{
+  THTensor *tensor = luaT_checkudata(L, 1, torch_Tensor);
+  THTensor_(retain)(cutorch_getstate(L), tensor);
+  return 0;
+}
+
 static int torch_Tensor_(free)(lua_State *L)
 {
   THTensor *tensor = luaT_checkudata(L, 1, torch_Tensor);
@@ -1112,6 +1119,8 @@ static int torch_Tensor_(read)(lua_State *L)
 }
 
 static const struct luaL_Reg torch_Tensor_(_) [] = {
+  {"retain", torch_Tensor_(retain)},
+  {"free", torch_Tensor_(free)},
   {"contiguous", torch_Tensor_(contiguous)},
   {"size", torch_Tensor_(size)},
   {"__len__", torch_Tensor_(size)},