Added ROCm support for nn_ops_test

2019-03-19 04:13:03 +00:00 · 2019-03-19 04:13:03 +00:00 · 4aaa257a5d
commit 4aaa257a5d
parent 949474a448
1 changed files with 2 additions and 2 deletions
--- a/tensorflow/core/kernels/nn_ops_test.cc
+++ b/tensorflow/core/kernels/nn_ops_test.cc
@ -15,9 +15,9 @@ limitations under the License.

 #define EIGEN_USE_THREADS

-#if GOOGLE_CUDA
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #define EIGEN_USE_GPU
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

 #include <functional>
 #include <memory>