diff options
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_module_group_metadata.cc | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc index 4f1715e4ca..bf33640db1 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc @@ -127,9 +127,14 @@ Status HloModuleGroupMetadata::VerifyCompanionSets() const { for (HloInstruction* instruction : *companions) { // Go through all the communicating instructions (send, recv) of the given // companion, and record their device. + auto it = tracked_instructions_comms_.find(instruction); + if (it == tracked_instructions_comms_.end()) { + // Companions can be added even if they have no communicating + // instructions, if they are parent of companions. + continue; + } std::unordered_set<int64> comm_devices; - for (HloInstruction* comm_instruction : - tracked_instructions_comms_.at(instruction)) { + for (HloInstruction* comm_instruction : it->second) { auto device = GetInstructionDevice(*comm_instruction); TF_RET_CHECK(device) << "Instruction " << comm_instruction->ToString() << " does not have a device"; |