Implemented cudaMemGetInfo for caching allocator (#600)

* Implemented cudaMemGetInfo for caching allocator

1 files changed, 3 insertions, 2 deletions

diff --git a/init.c b/init.c
index 69f5583..124be5c 100644
--- a/init.c
+++ b/init.c
@@ -698,13 +698,14 @@ static int cutorch_getMemoryUsage(lua_State *L) {
   size_t totalBytes = 0;
   int curDevice;
   THCudaCheck(cudaGetDevice(&curDevice));
+  THCState *state = cutorch_getstate(L);
 
   int device = luaL_optint(L, 1, -10);
   if (device == -10) { /* no argument passed, current device mem usage */
-    THCudaCheck(cudaMemGetInfo(&freeBytes, &totalBytes));
+    THCudaCheck(THCudaMemGetInfo(state, &freeBytes, &totalBytes));
   } else { /* argument was given, particular device's memory usage */
     THCudaCheck(cudaSetDevice(device-1)); /* zero indexed */
-    THCudaCheck(cudaMemGetInfo(&freeBytes, &totalBytes));
+    THCudaCheck(THCudaMemGetInfo(state, &freeBytes, &totalBytes));
     THCudaCheck(cudaSetDevice(curDevice));
   }
   lua_pushnumber(L, freeBytes);