Only update watermark when we had the correct batch size.
Only update watermark when we had the correct batch size.

PiperOrigin-RevId: 460243743
Change-Id: I4f9965340ad01a79bc24031848b099937ff5ce62

google/tcmalloc,google/tcmalloc,google/tcmalloc