experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Intel MKL] Fixing Device() API usage in mkl_eager_op_rewrite.cc

Browse Source 
API change in TF core caused a build failure in MKL CPU backend. This commit fixes the compilation failure.
This commit is contained in:
Niranjan Hasabnis 2020-02-20 16:52:54 -08:00
parent 7966f2df6b
commit 7fac3faba2
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorflow/core/common_runtime/eager/mkl_eager_op_rewrite.cc
View File

@ -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();
}