Call HloComputation.Accept instead of HloInstruction.Accept to get all instructions profiled.
RELNOTES: n/a PiperOrigin-RevId: 167640259
This commit is contained in:
parent
0ab137cd8b
commit
aebe8cc6f4
@ -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;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user