Commit 0ffc664b authored by Derek Murray's avatar Derek Murray Committed by TensorFlower Gardener
Browse files

[tf.data] Avoid tracking allocations when collecting stats for a `CapturedFunction`.

This change adds `NodeExecStatsInterface::TrackAllocations()` method to control allocation tracking, and uses it in the stats collector implementation used in `CapturedFunction`.

Currently, the use of any `StepStatsCollectorInterface` implementation will
enable allocation tracking. However, not all implementations actually make use
of the tracked allocations, and so this is wasted effort (and it has led to
memory leaks in the past). This change allows a `NodeExecStatsInterface` to
decide whether allocations should be tracked or not, on a node-by-node basis.
It also modifies the stats collector used in `tf.data` functions to disable
allocation tracking, since it is not (currently) used as part of the performance
modeling code.

PiperOrigin-RevId: 219847937
parent 0fc52d37
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment