Make class MultinomialLogisticLossLayerTest derive from CPUDeviceTest

AniChikage · May 26, 2015 · 68133e7 · 68133e7
1 parent 89bf3c3
commit 68133e7
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/caffe/test/test_multinomial_logistic_loss_layer.cpp b/src/caffe/test/test_multinomial_logistic_loss_layer.cpp
@@ -16,7 +16,7 @@
 namespace caffe {
 
 template <typename Dtype>
-class MultinomialLogisticLossLayerTest : public ::testing::Test {
+class MultinomialLogisticLossLayerTest : public CPUDeviceTest<Dtype> {
  protected:
   MultinomialLogisticLossLayerTest()
       : blob_bottom_data_(new Blob<Dtype>(10, 5, 1, 1)),
@@ -51,7 +51,6 @@ TYPED_TEST_CASE(MultinomialLogisticLossLayerTest, TestDtypes);
 
 TYPED_TEST(MultinomialLogisticLossLayerTest, TestGradientCPU) {
   LayerParameter layer_param;
-  Caffe::set_mode(Caffe::CPU);
   MultinomialLogisticLossLayer<TypeParam> layer(layer_param);
   layer.SetUp(this->blob_bottom_vec_, this->blob_top_vec_);
   GradientChecker<TypeParam> checker(1e-2, 2*1e-2, 1701, 0, 0.05);