Call HloComputation.Accept instead of HloInstruction.Accept to get all instructions profiled.

RELNOTES: n/a
PiperOrigin-RevId: 167640259
This commit is contained in:
A. Unique TensorFlower 2017-09-05 15:27:28 -07:00 committed by TensorFlower Gardener
parent 0ab137cd8b
commit aebe8cc6f4

View File

@ -198,8 +198,8 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault {
std::unordered_map<const HloInstruction*, size_t> hlo_to_profile_idx;
CollectProfileCandidates profile_candidates_for_computation(
&hlo_to_profile_idx);
TF_RETURN_IF_ERROR(computation->root_instruction()->Accept(
&profile_candidates_for_computation));
TF_RETURN_IF_ERROR(
computation->Accept(&profile_candidates_for_computation));
return hlo_to_profile_idx;
}