diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc index 9107a6456f2..8f494e53303 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc @@ -101,6 +101,7 @@ void CreateTPUBridgePipeline(OpPassManager &pm) { pm.addPass(TFDevice::CreateResourceOpLiftingPass()); pm.addPass(TF::CreateTFFunctionalControlFlowToRegions()); pm.addPass(mlir::createInlinerPass()); + pm.addPass(TFDevice::CreateMarkOpsForOutsideCompilationPass()); pm.addPass(CreateTPUExtractHeadTailOutsideCompilationPass()); pm.addPass(TF::CreateTFRegionControlFlowToFunctional());