diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc index 38c4da1f182..5d695b9c20f 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -315,6 +315,8 @@ Status HloComputation::RemoveInstructionImpl(HloInstruction* instruction, (*inst_it->second)->set_parent(nullptr); to_be_deleted_.emplace_back(inst_it->second->release()); to_be_deleted_.back()->DetachFromOperandsAndUsers(); + // Clear all operands to avoid Null operands. + to_be_deleted_.back()->RemoveAllOperands(); to_be_deleted_.back()->MarkAsDead(); instructions_.erase(inst_it->second); instruction_iterators_.erase(inst_it); diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index e21ae719e4d..9675a2f0f0d 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -1911,6 +1911,8 @@ class HloInstruction { // by factory methods. HloInstruction(HloOpcode opcode, const Shape& shape); + void RemoveAllOperands() { operands_.clear(); } + void RemoveOperandAt(int index) { operands_.erase(operands_.begin() + index); }