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

Added ROCm support for the extract_image_patches op

Browse Source
This commit is contained in:
Jeffrey Poznanovic 2019-03-19 18:12:09 +00:00
parent 3e98ebd9f5
commit 3fe83873bd
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/core/kernels/extract_image_patches_op.cc
View File

@ -130,7 +130,7 @@ TF_CALL_REAL_NUMBER_TYPES(REGISTER);
#undef REGISTER #undef REGISTER
#if GOOGLE_CUDA #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// Forward declarations of the functor specializations for GPU. // Forward declarations of the functor specializations for GPU.
namespace functor { namespace functor {
@ -160,6 +160,6 @@ TF_CALL_GPU_NUMBER_TYPES(REGISTER);
#undef REGISTER #undef REGISTER
#endif // GOOGLE_CUDA #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow } // namespace tensorflow

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

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