Automated rollback of commit 13f0092fa3

PiperOrigin-RevId: 294292388
Change-Id: Ie74a6c898b9d8866bb89644f5b19bc9d7867a8a1
This commit is contained in:
A. Unique TensorFlower 2020-02-10 13:35:55 -08:00 committed by TensorFlower Gardener
parent 7343cfb70a
commit 036b77ad19

View File

@ -141,9 +141,12 @@ mlir::OwningModuleRef SavedModelV1ToMlirImport(
absl::string_view saved_model_dir,
const std::unordered_set<std::string>& tags, mlir::MLIRContext* context) {
tensorflow::SavedModelBundle bundle;
auto load_status = tensorflow::LoadSavedModel(
/* session_options = */ {}, /* run_options = */ {},
std::string(saved_model_dir), tags, &bundle);
tensorflow::SessionOptions session_options;
// Force saved model states to be restored to CPU.
(*session_options.config.mutable_device_count())["GPU"] = 0;
auto load_status =
tensorflow::LoadSavedModel(session_options, /* run_options = */ {},
std::string(saved_model_dir), tags, &bundle);
if (!load_status.ok()) {
LOG(ERROR) << "Failed to load saved model v1 '" << saved_model_dir
<< "': " << load_status;