Make class ArgMaxLayerTest derive from CPUDeviceTest

diff --git a/src/caffe/test/test_argmax_layer.cpp b/src/caffe/test/test_argmax_layer.cpp
index 3487d42..895c3d3 100644 (file)
--- a/src/caffe/test/test_argmax_layer.cpp
+++ b/src/caffe/test/test_argmax_layer.cpp
@@ -13,13 +13,12 @@
 namespace caffe {
 
 template <typename Dtype>
-class ArgMaxLayerTest : public ::testing::Test {
+class ArgMaxLayerTest : public CPUDeviceTest<Dtype> {
  protected:
   ArgMaxLayerTest()
       : blob_bottom_(new Blob<Dtype>(10, 20, 1, 1)),
         blob_top_(new Blob<Dtype>()),
         top_k_(5) {
-    Caffe::set_mode(Caffe::CPU);
     Caffe::set_random_seed(1701);
     // fill the values
     FillerParameter filler_param;