VirtualCluster constructor takes a node_estimator pointer.

PiperOrigin-RevId: 159289065

2 files changed, 8 insertions, 0 deletions

diff --git a/tensorflow/core/grappler/clusters/virtual_cluster.cc b/tensorflow/core/grappler/clusters/virtual_cluster.cc
index 30eaccb801..32d9be9aee 100644
--- a/tensorflow/core/grappler/clusters/virtual_cluster.cc
+++ b/tensorflow/core/grappler/clusters/virtual_cluster.cc
@@ -29,6 +29,12 @@ VirtualCluster::VirtualCluster(
   devices_ = devices;
 }
 
+VirtualCluster::VirtualCluster(
+    const std::unordered_map<string, DeviceProperties>& devices,
+    OpLevelCostEstimator* node_estimator)
+    : Cluster(0), node_estimator_(node_estimator) {
+  devices_ = devices;
+}
 VirtualCluster::~VirtualCluster() {}
 
 Status VirtualCluster::Provision() { return Status::OK(); }
diff --git a/tensorflow/core/grappler/clusters/virtual_cluster.h b/tensorflow/core/grappler/clusters/virtual_cluster.h
index 452c6f7c8e..a74911cb23 100644
--- a/tensorflow/core/grappler/clusters/virtual_cluster.h
+++ b/tensorflow/core/grappler/clusters/virtual_cluster.h
@@ -30,6 +30,8 @@ namespace grappler {
 class VirtualCluster : public Cluster {
  public:
   VirtualCluster(const std::unordered_map<string, DeviceProperties>& devices);
+  VirtualCluster(const std::unordered_map<string, DeviceProperties>& devices,
+                 OpLevelCostEstimator* node_estimator);
 
   ~VirtualCluster() override;