message: [TF2XLA] Use Two pass argmax on CPU as well for CatagoricalOp.
diff: @@ -103,14 +103,8 @@ class CategoricalOp : public XlaOpKernel {
  xla::PrimitiveType xla_output_type;
  OP_REQUIRES_OK(ctx,
  DataTypeToPrimitiveType(output_type(0), &xla_output_type));
  - xla::XlaOp argmax;
  - if (is_gpu_) {
  - argmax = xla::ArgMaxTwoPass(softmax_entries, xla_output_type,
  + xla::XlaOp argmax = xla::ArgMaxTwoPass(softmax_entries, xla_output_type,
  /*axis=*/class_dimension);
  - } else {
  - argmax = xla::ArgMax(softmax_entries, xla_output_type,
  - /*axis=*/class_dimension);
  - }
  if (num_samples == 1) {
  argmax = xla::Reshape(argmax, {batch_size, 1});
  