experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #26742 from ROCmSoftwarePlatform:google_upstream_argmax_op

Browse Source 
PiperOrigin-RevId: 239163567
This commit is contained in:
TensorFlower Gardener 2019-03-19 04:22:48 -07:00
commit 83c6c627e5
2 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorflow/core/kernels/argmax_op.cc
View File

@ -17,9 +17,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 "tensorflow/core/kernels/argmax_op.h"
@ -150,7 +150,7 @@ class ArgMinOp
TF_CALL_REAL_NUMBER_TYPES(REGISTER_ARGMAX);
#if GOOGLE_CUDA
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// Forward declarations of the functor specializations for GPU.
namespace functor {
@ -226,6 +226,6 @@ TF_CALL_GPU_NUMBER_TYPES(REGISTER_ARGMAX_GPU);
#undef REGISTER_ARGMAX_GPU
#endif // GOOGLE_CUDA
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow

4
tensorflow/core/kernels/argmax_op_gpu.cu.cc
View File

@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#if GOOGLE_CUDA
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
@ -34,4 +34,4 @@ TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_SPEC);
} // end namespace tensorflow
#endif // GOOGLE_CUDA
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM