From d28f7e6681ffe02a151b9dc89098d7fcef50b214 Mon Sep 17 00:00:00 2001
From: Alexey <AlexeyAB@users.noreply.github.com>
Date: Wed, 28 Mar 2018 20:51:14 +0000
Subject: [PATCH] Update Readme.md

---
 src/yolo_v2_class.cpp |   14 +++++++++-----
 1 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/yolo_v2_class.cpp b/src/yolo_v2_class.cpp
index c805e29..076bab8 100644
--- a/src/yolo_v2_class.cpp
+++ b/src/yolo_v2_class.cpp
@@ -22,12 +22,14 @@
 
 #define FRAMES 3
 
+#ifdef GPU
 void check_cuda(cudaError_t status) {
 	if (status != cudaSuccess) {
 		const char *s = cudaGetErrorString(status);
 		printf("CUDA Error Prev: %s\n", s);
 	}
 }
+#endif
 
 struct detector_gpu_t {
 	float **probs;
@@ -52,11 +54,12 @@
 	detector_gpu_t &detector_gpu = *static_cast<detector_gpu_t *>(detector_gpu_ptr.get());
 
 #ifdef GPU
-	check_cuda( cudaSetDevice(gpu_id) );
-	printf(" Used GPU %d \n", gpu_id);
+	//check_cuda( cudaSetDevice(cur_gpu_id) );
+	cuda_set_device(cur_gpu_id);
+	printf(" Used GPU %d \n", cur_gpu_id);
 #endif
 	network &net = detector_gpu.net;
-	net.gpu_index = gpu_id;
+	net.gpu_index = cur_gpu_id;
 	//gpu_index = i;
 	
 	char *cfgfile = const_cast<char *>(cfg_filename.data());
@@ -67,7 +70,7 @@
 		load_weights(&net, weightfile);
 	}
 	set_batch_network(&net, 1);
-	net.gpu_index = gpu_id;
+	net.gpu_index = cur_gpu_id;
 
 	layer l = net.layers[net.n - 1];
 	int j;
@@ -107,7 +110,8 @@
 	int old_gpu_index;
 #ifdef GPU
 	cudaGetDevice(&old_gpu_index);
-	cudaSetDevice(detector_gpu.net.gpu_index);
+	//cudaSetDevice(detector_gpu.net.gpu_index);
+	cuda_set_device(detector_gpu.net.gpu_index);
 #endif
 
 	free_network(detector_gpu.net);

--
Gitblit v1.10.0