Adding GPU coverage to the DataLayer test.

1 files changed, 16 insertions, 0 deletions

diff --git a/src/caffe/test/test_data_layer.cpp b/src/caffe/test/test_data_layer.cpp
index 66e99568..719c50eb 100644
--- a/src/caffe/test/test_data_layer.cpp
+++ b/src/caffe/test/test_data_layer.cpp
@@ -81,6 +81,7 @@ TYPED_TEST(DataLayerTest, TestRead) {
   EXPECT_EQ(this->blob_top_label_->channels(), 1);
   EXPECT_EQ(this->blob_top_label_->height(), 1);
   EXPECT_EQ(this->blob_top_label_->width(), 1);
+
   // Go through the data 100 times
   for (int iter = 0; iter < 100; ++iter) {
     layer.Forward(this->blob_bottom_vec_, &this->blob_top_vec_);
@@ -94,6 +95,21 @@ TYPED_TEST(DataLayerTest, TestRead) {
       }
     }
   }
+
+  // Same test, in GPU mode.
+  Caffe::set_mode(Caffe::GPU);
+  for (int iter = 0; iter < 100; ++iter) {
+    layer.Forward(this->blob_bottom_vec_, &this->blob_top_vec_);
+    for (int i = 0; i < 5; ++i) {
+      EXPECT_EQ(i, this->blob_top_label_->cpu_data()[i]);
+    }
+    for (int i = 0; i < 5; ++i) {
+      for (int j = 0; j < 24; ++j) {
+        EXPECT_EQ(i, this->blob_top_data_->cpu_data()[i * 24 + j])
+            << "debug: i " << i << " j " << j;
+      }
+    }
+  }
 }
 
 }