Searched refs:computation_peak_memory_ (Results 1 – 2 of 2) sorted by relevance
965 TF_RET_CHECK(ContainsKey(computation_peak_memory_, computation)); in CalledComputationsMemoryUsage()966 callee_usage += computation_peak_memory_.at(computation); in CalledComputationsMemoryUsage()977 << HumanReadableNumBytes(computation_peak_memory_.at(computation)); in RematerializeComputation()1181 << HumanReadableNumBytes(computation_peak_memory_.at(computation)) in RematerializeComputation()1185 computation_peak_memory_.at(computation) = peak_memory; in RematerializeComputation()1209 computation_peak_memory_.clear(); in Run()1242 computation_peak_memory_[node.computation()], in Run()1255 computation_peak_memory_.at(module->entry_computation()) + in Run()1279 computation_peak_memory_.at(module->entry_computation()) + in Run()
119 absl::flat_hash_map<const HloComputation*, int64> computation_peak_memory_; variable