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

Merge pull request #36944 from Intel-tensorflow:nhasabni/mkl_eager_op_rewrite_fix_v2

Browse Source 
PiperOrigin-RevId: 296386102
Change-Id: I4ef2f4d104011ff01ce050eac7041827d2c907b6
This commit is contained in:
TensorFlower Gardener 2020-02-21 00:58:41 -08:00
commit 386ffb44ab
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();
}