TF_RETURN_IF_ERROR(ops::SparseSoftmaxCrossEntropyWithLogits(
      ctx, inputs[0], inputs[1], &loss, &backprop,
      "SparseSoftmaxCrossEntropyWithLogits"));
  // `gradient_checker` only works with model that returns only 1 tensor.
  // Although, `ops::SparseSoftmaxCrossEntropyWithLogits` returns 2 tensors, the
  // second tensor isn't needed for computing gradient so we could safely drop
  // it.
  outputs[0] = loss;