diff --git a/client.cu b/client.cu
index 85ca7bd..5762e18 100644
--- a/client.cu
+++ b/client.cu
@@ -145,10 +145,10 @@ nvmlReturn_t nvmlInit_v2()
     return send_rpc_message(RPC_nvmlInit_v2);
 }
 
-nvmlReturn_t nvmlDeviceGetCountHandler_v2()
+nvmlReturn_t nvmlDeviceGetCountHandler_v2(unsigned int* device_count)
 {
     open_rpc_client();
-    return send_rpc_message(RPC_nvmlDeviceGetCount_v2);
+    return send_rpc_message(RPC_nvmlDeviceGetCount_v2, {}, {{device_count, sizeof(unsigned int)}});
 }
 
 nvmlReturn_t nvmlShutdown()
diff --git a/server.cu b/server.cu
index fbc6fd3..76fb914 100644
--- a/server.cu
+++ b/server.cu
@@ -82,9 +82,9 @@ void *client_handler(void *arg)
                 printf("Number of devices: %u\n", dcount);
             }
 
-            if (write(connfd, &result, sizeof(nvmlReturn_t)) < 0)
-                goto exit;
-            }
+            if (write(connfd, &dcount, sizeof(unsigned int)) <= 0)
+	            goto exit;
+	        }
         }
     }