Fixed memory release

8 years ago · 576dbe12e6
parent 51477ab274
commit 576dbe12e6
1 changed files with 4 additions and 1 deletions
--- a/src/network.c
+++ b/src/network.c
@ -595,11 +595,14 @@ void free_network(network net)
 		free_layer(net.layers[i]);
 	}
 	free(net.layers);
-	free(net.workspace);
 #ifdef GPU
+	if (gpu_index >= 0) cuda_free(net.workspace);
+	else free(net.workspace);
 	if (*net.input_gpu) cuda_free(*net.input_gpu);
 	if (*net.truth_gpu) cuda_free(*net.truth_gpu);
 	if (net.input_gpu) free(net.input_gpu);
 	if (net.truth_gpu) free(net.truth_gpu);
+#else
+	free(net.workspace);
 #endif
 }