diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
index 080e246458e..50c9e64deb6 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include "cuda/include/cuda_runtime_api.h"
 #include "tensorflow/contrib/tensorrt/log/trt_logger.h"
 #include "tensorflow/core/platform/logging.h"
-//#include "tensorflow/core/framework/device_base.h"
 #include "tensorflow/core/platform/stream_executor.h"
 
 namespace tensorflow {
diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/contrib/tensorrt/segment/segment.cc
index c9d3840606b..16722266435 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.cc
+++ b/tensorflow/contrib/tensorrt/segment/segment.cc
@@ -223,15 +223,6 @@ tensorflow::Status SegmentGraph(
     }
   }
 
-  // Cleanup the graph to remove disconnected nodes before outputting
-  if (VLOG_IS_ON(2)) {
-    for (tensorflow::Node* node : graph.nodes()) {
-      if ((node->in_edges().size() == 0) && (node->out_edges().size() == 0)) {
-        graph.RemoveNode(node);
-      }
-    }
-  }
-
   // Convert the segments into the expected return format
   for (const auto& itr : sg_map) {
     const auto& segment_node_names = itr.second;