Get rid of to_int64 deprecation warning in the logs.

PiperOrigin-RevId: 225256193

tensorflow/python/ops/math_ops.py

@@ -1465,7 +1465,7 @@ def count_nonzero_v2(input,  # pylint: disable=redefined-builtin
     return cast(
         reduce_sum(
             # int64 reduction happens on GPU
-            to_int64(gen_math_ops.not_equal(input, zero)),
+            cast(gen_math_ops.not_equal(input, zero), dtypes.int64),
             axis=axis,
             keepdims=keepdims),
         dtype=dtype)