From 036b77ad19da6e5c43009c5df010c2a65710107f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 10 Feb 2020 13:35:55 -0800
Subject: [PATCH] Automated rollback of commit
 13f0092fa38f9c9753f60c90a09aee8c9c9751db

PiperOrigin-RevId: 294292388
Change-Id: Ie74a6c898b9d8866bb89644f5b19bc9d7867a8a1
---
 .../mlir/tensorflow/translate/tf_mlir_translate.cc       | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/translate/tf_mlir_translate.cc b/tensorflow/compiler/mlir/tensorflow/translate/tf_mlir_translate.cc
index b4b5b869e74..5e958960d07 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/tf_mlir_translate.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/tf_mlir_translate.cc
@@ -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;