Avoid using cudaMemcpy for memcpy when there is no GPU and CUDA driver

diff --git a/src/caffe/test/test_math_functions.cpp b/src/caffe/test/test_math_functions.cpp
index 941d8b9..ddb9f06 100644 (file)
--- a/src/caffe/test/test_math_functions.cpp
+++ b/src/caffe/test/test_math_functions.cpp
@@ -209,6 +209,7 @@ TYPED_TEST(MathFunctionsTest, TestCopyCPU) {
   const int n = this->blob_bottom_->count();
   const TypeParam* bottom_data = this->blob_bottom_->cpu_data();
   TypeParam* top_data = this->blob_top_->mutable_cpu_data();
+  Caffe::set_mode(Caffe::CPU);
   caffe_copy(n, bottom_data, top_data);
   for (int i = 0; i < n; ++i) {
     EXPECT_EQ(bottom_data[i], top_data[i]);
@@ -219,6 +220,7 @@ TYPED_TEST(MathFunctionsTest, TestCopyGPU) {
   const int n = this->blob_bottom_->count();
   const TypeParam* bottom_data = this->blob_bottom_->gpu_data();
   TypeParam* top_data = this->blob_top_->mutable_gpu_data();
+  Caffe::set_mode(Caffe::GPU);
   caffe_copy(n, bottom_data, top_data);
   bottom_data = this->blob_bottom_->cpu_data();
   top_data = this->blob_top_->mutable_cpu_data();

diff --git a/src/caffe/util/math_functions.cpp b/src/caffe/util/math_functions.cpp
index 918bb3c..df32093 100644 (file)
--- a/src/caffe/util/math_functions.cpp
+++ b/src/caffe/util/math_functions.cpp
@@ -152,7 +152,11 @@ void caffe_add_scalar(const int N, const double alpha, double* Y) {
 template <typename Dtype>
 void caffe_copy(const int N, const Dtype* X, Dtype* Y) {
   if (X != Y) {
-    CUDA_CHECK(cudaMemcpy(Y, X, sizeof(Dtype) * N, cudaMemcpyDefault));
+    if (Caffe::mode() == Caffe::CPU) {
+      memcpy(Y, X, sizeof(Dtype) * N);
+    } else {
+      CUDA_CHECK(cudaMemcpy(Y, X, sizeof(Dtype) * N, cudaMemcpyDefault));
+    }
   }
 }
 
@@ -164,7 +168,11 @@ template void caffe_copy<double>(const int N, const double* X, double* Y);
 
 void caffe_memcpy(const size_t N, const void* X, void* Y) {
   if (X != Y) {
-    CUDA_CHECK(cudaMemcpy(Y, X, N, cudaMemcpyDefault));
+    if (Caffe::mode() == Caffe::CPU) {
+      memcpy(Y, X, N);
+    } else {
+      CUDA_CHECK(cudaMemcpy(Y, X, N, cudaMemcpyDefault));
+    }
   }
 }