Memory leak fix

5 years ago · 03eee3092f
parent 594f7ce6af
commit 03eee3092f
2 changed files with 2 additions and 2 deletions
--- a/src/layer.c
+++ b/src/layer.c
@ -193,7 +193,7 @@ void free_layer_custom(layer l, int keep_cudnn_desc)
        if (l.output_avg_gpu)          cuda_free(l.output_avg_gpu), l.output_avg_gpu = NULL;
        if (l.activation_input_gpu)    cuda_free(l.activation_input_gpu), l.activation_input_gpu = NULL;
    }
-    if (l.delta_gpu && l.keep_delta_gpu && l.optimized_memory < 3) cuda_free(l.delta_gpu), l.delta_gpu = NULL;
+    if (l.delta_gpu && (l.optimized_memory < 1 || l.keep_delta_gpu && l.optimized_memory < 3)) cuda_free(l.delta_gpu), l.delta_gpu = NULL;
    if (l.rand_gpu)                cuda_free(l.rand_gpu);
    if (l.squared_gpu)             cuda_free(l.squared_gpu);
    if (l.norms_gpu)               cuda_free(l.norms_gpu);
--- a/src/parser.c
+++ b/src/parser.c
@ -1457,9 +1457,9 @@ network parse_network_cfg_custom(char *filename, int batch, int time_steps)

 #ifdef GPU
        // futher GPU-memory optimization: net.optimized_memory == 2
+        l.optimized_memory = net.optimized_memory;
        if (net.optimized_memory >= 2 && params.train && l.type != DROPOUT)
        {
-            l.optimized_memory = net.optimized_memory;
            if (l.output_gpu) {
                cuda_free(l.output_gpu);
                //l.output_gpu = cuda_make_array_pinned(l.output, l.batch*l.outputs); // l.steps