From 7fac3faba23fbe39f47fa2ae1add788144d12fa6 Mon Sep 17 00:00:00 2001
From: Niranjan Hasabnis <niranjan.hasabnis@intel.com>
Date: Thu, 20 Feb 2020 16:52:54 -0800
Subject: [PATCH] [Intel MKL] Fixing Device() API usage in
 mkl_eager_op_rewrite.cc

API change in TF core caused a build failure in MKL CPU backend. This commit fixes the compilation failure.
---
 .../core/common_runtime/eager/mkl_eager_op_rewrite.cc     | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/common_runtime/eager/mkl_eager_op_rewrite.cc b/tensorflow/core/common_runtime/eager/mkl_eager_op_rewrite.cc
index a222724ec8f..346ccc11ca8 100644
--- a/tensorflow/core/common_runtime/eager/mkl_eager_op_rewrite.cc
+++ b/tensorflow/core/common_runtime/eager/mkl_eager_op_rewrite.cc
@@ -135,11 +135,13 @@ Status MklEagerOpRewrite::SetupNewOp(
       ->MutableAttrs()
       ->Set("_kernel", mkl_op_registry::kMklNameChangeOpLabel);
 
-  if (orig_op->Device() != nullptr) {
-    (*new_mkl_op)->SetDevice(orig_op->Device());
-  } else {
+  if (orig_op->Device() == kVariantDeviceNull) {
     string device_name = orig_op->GetDeviceName();
     (*new_mkl_op)->SetDeviceName(device_name.c_str());
+  } else if (VariantDeviceIsCustom(orig_op->Device())) {
+    (*new_mkl_op)->SetDevice(absl::get<CustomDevice*>(orig_op->Device()));
+  } else {
+    (*new_mkl_op)->SetDevice(absl::get<Device*>(orig_op->Device()));
   }
   return Status::OK();
 }