Merge pull request tensorflow#1503 from ROCmSoftwarePlatform/deven/av…

…gpooling_unit_test_fix

 Fix for a regression in the unit test `//tensorflow/python/kernel_tests/nn_ops:pooling_ops_test_gpu`

tensorflow/core/kernels/eigen_pooling.h

@@ -284,7 +284,11 @@ struct AvgPoolMeanReducer {
 
   EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE AvgPoolMeanReducer() : scalarCount_(0) {
     typedef typename packet_traits<T>::type Packet;
+#if defined(__HIPCC__)
+    packetCount_ = 0;
+#else
     packetCount_ = pset1<Packet>(T(0.0));
+#endif
   }
 
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) {
@@ -305,6 +309,7 @@ struct AvgPoolMeanReducer {
 
 #if (EIGEN_ARCH_i386 || EIGEN_ARCH_x86_64) && !defined(__CUDACC__) && \
     !defined(__HIPCC__)
+
 #ifdef EIGEN_VECTORIZE_AVX512
 #define pequal(a, b)   \
   _mm512_castsi512_ps( \
@@ -364,7 +369,11 @@ struct AvgPoolMeanReducer {
  protected:
   typedef typename packet_traits<T>::type Packet;
   int scalarCount_;
+#if defined(__HIPCC__)
+  int packetCount_;
+#else
   Packet packetCount_;
+#endif
 };
 
 template <typename Device>

tensorflow/python/kernel_tests/nn_ops/BUILD

@@ -496,7 +496,7 @@ cuda_py_test(
     srcs = ["pooling_ops_test.py"],
     shard_count = 10,
     # Some operations in this test can only be checked on sm61+.
-    tags = ["prefer-sm70", "no_rocm"],
+    tags = ["prefer-sm70"],
     deps = [
         "//tensorflow/python:array_ops",
         "//tensorflow/python:client_testlib",