From ec3d050a76ee8c41f35c4531d3fa07a2d9c28ed3 Mon Sep 17 00:00:00 2001
From: Joseph Redmon <pjreddie@gmail.com>
Date: Thu, 02 Jun 2016 22:25:24 +0000
Subject: [PATCH] hope i didn't break anything
---
src/network.c | 10 ++++++++++
1 files changed, 10 insertions(+), 0 deletions(-)
diff --git a/src/network.c b/src/network.c
index 8f39f7b..88b7085 100644
--- a/src/network.c
+++ b/src/network.c
@@ -137,6 +137,7 @@
void forward_network(network net, network_state state)
{
+ state.workspace = net.workspace;
int i;
for(i = 0; i < net.n; ++i){
state.index = i;
@@ -400,6 +401,7 @@
net->w = w;
net->h = h;
int inputs = 0;
+ size_t workspace_size = 0;
//fprintf(stderr, "Resizing to %d x %d...", w, h);
//fflush(stderr);
for (i = 0; i < net->n; ++i){
@@ -419,12 +421,20 @@
}else{
error("Cannot resize this type of layer");
}
+ if(l.workspace_size > workspace_size) workspace_size = l.workspace_size;
inputs = l.outputs;
net->layers[i] = l;
w = l.out_w;
h = l.out_h;
if(l.type == AVGPOOL) break;
}
+#ifdef GPU
+ cuda_free(net->workspace);
+ net->workspace = cuda_make_array(0, (workspace_size-1)/sizeof(float)+1);
+#else
+ free(net->workspace);
+ net->workspace = calloc(1, (workspace_size-1)/sizeof(float)+1);
+#endif
//fprintf(stderr, " Done!\n");
return 0;
}
--
Gitblit v1.10.0