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; std::unordered_map<const HloInstruction*, size_t> hlo_to_profile_idx;
CollectProfileCandidates profile_candidates_for_computation( CollectProfileCandidates profile_candidates_for_computation(
&hlo_to_profile_idx); &hlo_to_profile_idx);
TF_RETURN_IF_ERROR(computation->root_instruction()->Accept( TF_RETURN_IF_ERROR(
&profile_candidates_for_computation)); computation->Accept(&profile_candidates_for_computation));
return hlo_to_profile_idx; return hlo_to_profile_idx;
} }