Joseph Redmon
2016-08-05 845ab7579685b6702c92c1088ec11e71bde51f3c
src/convolutional_layer.c
@@ -104,6 +104,7 @@
size_t get_workspace_size(layer l){
#ifdef CUDNN
    if(gpu_index >= 0){
    size_t most = 0;
    size_t s = 0;
    cudnnGetConvolutionForwardWorkspaceSize(cudnn_handle(),
@@ -131,9 +132,9 @@
            &s);
    if (s > most) most = s;
    return most;
#else
    return (size_t)l.out_h*l.out_w*l.size*l.size*l.c*sizeof(float);
    }
#endif
    return (size_t)l.out_h*l.out_w*l.size*l.size*l.c*sizeof(float);
}
#ifdef GPU
@@ -240,6 +241,7 @@
    }
#ifdef GPU
    if(gpu_index >= 0){
    l.filters_gpu = cuda_make_array(l.filters, c*n*size*size);
    l.filter_updates_gpu = cuda_make_array(l.filter_updates, c*n*size*size);
@@ -283,6 +285,7 @@
    cudnnCreateConvolutionDescriptor(&l.convDesc);
    cudnn_convolutional_setup(&l);
#endif
    }
#endif
    l.workspace_size = get_workspace_size(l);
    l.activation = activation;