From 28a637415e288f71f23a4006e99767623e0294b8 Mon Sep 17 00:00:00 2001
From: YuanRisheng <yuanrisheng@baidu.com>
Date: Mon, 18 Oct 2021 20:23:30 +0800
Subject: [PATCH] Perfect unitests (#20)

* perfect unittest

* update license

* fix bug when run tcmpt_utils_test
---
 paddle/fluid/framework/tcmpt_utils_test.cc | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/paddle/fluid/framework/tcmpt_utils_test.cc b/paddle/fluid/framework/tcmpt_utils_test.cc
index c5af18f6f65aa5..f1966789c1ddee 100644
--- a/paddle/fluid/framework/tcmpt_utils_test.cc
+++ b/paddle/fluid/framework/tcmpt_utils_test.cc
@@ -49,12 +49,17 @@ TEST(TcmptUtils, VarToPtTensor) {
   auto* data =
       value->mutable_data<int>(make_ddim({1, 1}), paddle::platform::CPUPlace());
   data[0] = 123;
-  auto tensor_def = pt::TensorArgDef(pt::Backend::kCUDA, pt::DataLayout::kNCHW,
+  pt::Backend expect_backend = pt::Backend::kCPU;
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  expect_backend = pt::Backend::kCUDA;
+#endif
+  auto tensor_def = pt::TensorArgDef(expect_backend, pt::DataLayout::kNCHW,
                                      pt::DataType::kINT32);
   // 2. test API
   auto tensor_x = InputVariableToPtTensor(v, tensor_def);
   // 3. check result
-  ASSERT_EQ(tensor_x->backend(), pt::Backend::kCUDA);
+  ASSERT_EQ(tensor_x->backend(), expect_backend);
   ASSERT_EQ(tensor_x->type(), pt::DataType::kINT32);
 }