diff --git a/src/caffe/parallel.cpp b/src/caffe/parallel.cpp index 6e7d802bb99..d48136c56b8 100644 --- a/src/caffe/parallel.cpp +++ b/src/caffe/parallel.cpp @@ -64,7 +64,9 @@ static size_t total_size(const vector*>& params) { size_t size = 0; for (int i = 0; i < params.size(); ++i) size += params[i]->count(); - return size; + // Size have at least one byte, otherwise cudaMalloc fails if net has no + // learnable parameters. + return (size > 0) ? size : 1; } template