Replace `OpKernelContext::wrapped_allocators()` with `ConsumeWrappedAllocators()`.
This change makes the ownership transfer from `OpKernelContext` the stats collection machinery explicit, and should prevent memory leaks in future. It also logs a warning if the wrapped allocators are not consumed by a stats collector, which indicates that the code is taking a more expensive allocation path for no purpose. PiperOrigin-RevId: 220171712
Loading
Please sign in to comment