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