message: Consider all sub-graphs for tensor allocation tracking
diff: @@ -152,10 +152,10 @@ AllocationFinder::FindConsumers(HloInstruction* inst) {
  }
  Status AllocationFinder::CreateAllocationMap(HloModule* module) {
  - FindAllocatingInstructions finder;
  + for (const auto& comp : module->computations()) {
  - auto entry = module->entry_computation();
  - TF_RETURN_IF_ERROR(entry->Accept(&finder));
  + FindAllocatingInstructions finder;
  + TF_RETURN_IF_ERROR(comp->Accept(&finder));
  for (auto inst : finder.allocating_instructions) {
  TensorTarget target = FindConsumers(inst);
  @@ -163,6 +163,7 @@ Status AllocationFinder::CreateAllocationMap(HloModule* module) {
  tensor_allocation_map.insert(std::make_pair(inst, target));
  }
  }
  + }
  return Status::OK();
  }
  