diff --git a/lib/TH/THGeneral.c b/lib/TH/THGeneral.c index c838108f..3b9ed15c 100644 --- a/lib/TH/THGeneral.c +++ b/lib/TH/THGeneral.c @@ -268,3 +268,8 @@ double THLog1p(const double x) return log1p(x); #endif } + +long THGetHeapSize() +{ + return heapSize; +} diff --git a/lib/TH/THGeneral.h.in b/lib/TH/THGeneral.h.in index 3b62e4a0..b306072d 100644 --- a/lib/TH/THGeneral.h.in +++ b/lib/TH/THGeneral.h.in @@ -51,6 +51,7 @@ TH_API void* THAlloc(long size); TH_API void* THRealloc(void *ptr, long size); TH_API void THFree(void *ptr); TH_API void THSetGCHandler( void (*torchGCHandlerFunction)(void *data), void *data ); +TH_API long THGetHeapSize(void); // this hook should only be called by custom allocator functions TH_API void THHeapUpdate(long size); diff --git a/utils.c b/utils.c index 35fdae40..bdd7996e 100644 --- a/utils.c +++ b/utils.c @@ -227,6 +227,12 @@ static int torch_updateerrorhandlers(lua_State *L) return 0; } +static int torch_getmemory(lua_State *L) +{ + lua_pushinteger(L, THGetHeapSize()); + return 1; +} + static const struct luaL_Reg torch_utils__ [] = { {"getdefaulttensortype", torch_lua_getdefaulttensortype}, {"isatty", torch_isatty}, @@ -250,6 +256,7 @@ static const struct luaL_Reg torch_utils__ [] = { {"pointer", luaT_lua_pointer}, {"setheaptracking", torch_setheaptracking}, {"updateerrorhandlers", torch_updateerrorhandlers}, + {"getmemory", torch_getmemory}, {NULL, NULL} };