diff options
Diffstat (limited to 'tensorflow/core/grappler')
| -rw-r--r-- | tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc | 8 |
1 files changed, 8 insertions, 0 deletions
diff --git a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc index 8a65cd3ec3..c8f9311b2e 100644 --- a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc @@ -156,6 +156,9 @@ string TryFindHostDevice(const gtl::FlatSet<string>& devices, // We couldn't find an appropriate Host device, return original device. return device; } + +// All the nodes that should be blacklisted and not swapped. +bool IsBlacklisted(const NodeDef& node) { return IsCollective(node); } } // end namespace internal Status PinToHostOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item, @@ -185,6 +188,11 @@ Status PinToHostOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item, continue; } + // Skip these node types. + if (internal::IsBlacklisted(node)) { + continue; + } + // Check the node can be run on CPU. Status s = FindKernelDef(DEVICE_CPU, node, nullptr, nullptr); if (!s.ok()) { |