Make class CuDNNSoftmaxLayerTest derive from GPUDeviceTest

diff --git a/src/caffe/test/test_softmax_layer.cpp b/src/caffe/test/test_softmax_layer.cpp
index f667442..996da4b 100644 (file)
--- a/src/caffe/test/test_softmax_layer.cpp
+++ b/src/caffe/test/test_softmax_layer.cpp
@@ -82,7 +82,7 @@ TYPED_TEST(SoftmaxLayerTest, TestGradient) {
 
 #ifdef USE_CUDNN
 template <typename Dtype>
-class CuDNNSoftmaxLayerTest : public ::testing::Test {
+class CuDNNSoftmaxLayerTest : public GPUDeviceTest<Dtype> {
  protected:
   CuDNNSoftmaxLayerTest()
       : blob_bottom_(new Blob<Dtype>(2, 10, 2, 3)),
@@ -104,7 +104,6 @@ class CuDNNSoftmaxLayerTest : public ::testing::Test {
 TYPED_TEST_CASE(CuDNNSoftmaxLayerTest, TestDtypes);
 
 TYPED_TEST(CuDNNSoftmaxLayerTest, TestForwardCuDNN) {
-  Caffe::set_mode(Caffe::GPU);
   LayerParameter layer_param;
   CuDNNSoftmaxLayer<TypeParam> layer(layer_param);
   layer.SetUp(this->blob_bottom_vec_, this->blob_top_vec_);
@@ -138,7 +137,6 @@ TYPED_TEST(CuDNNSoftmaxLayerTest, TestForwardCuDNN) {
 }
 
 TYPED_TEST(CuDNNSoftmaxLayerTest, TestGradientCuDNN) {
-  Caffe::set_mode(Caffe::GPU);
   LayerParameter layer_param;
   CuDNNSoftmaxLayer<TypeParam> layer(layer_param);
   GradientChecker<TypeParam> checker(1e-2, 1e-3);