Fix ROCM build after #39429 and ce2a22c1dd

ROCM does not know about `struct NumericTraits<>` template, so we need to add back the conditional define block.

PiperOrigin-RevId: 337963975
Change-Id: Ie8014a51e0e50e1808db2c4cd4c49b0b83ee82cb

tensorflow/core/kernels/topk_op_gpu.h

@@ -35,12 +35,14 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/gpu_kernel_helper.h"
 
+#if GOOGLE_CUDA
 // Required for sorting Eigen::half
 namespace cub {
 template <>
 struct NumericTraits<Eigen::half>
     : BaseTraits<FLOATING_POINT, true, false, uint16_t, Eigen::half> {};
 }  // namespace cub
+#endif  // GOOGLE_CUDA
 
 namespace tensorflow {