Explicitly specify the value of all the attributes for the newly created Assign

nodes since we can't always rely on TF calling AddDefaultAttrsToGraphDef.

PiperOrigin-RevId: 184698463

diff --git a/tensorflow/core/grappler/optimizers/memory_optimizer.cc b/tensorflow/core/grappler/optimizers/memory_optimizer.cc
index be18e7d6f01ca4c5cbabb9099685fa7b8efeb1e3..9f3e94052f8289f575abdc6326cf2ee31749a8e6 100644 (file)
--- a/tensorflow/core/grappler/optimizers/memory_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/memory_optimizer.cc
@@ -641,6 +641,8 @@ bool SchedulingPass(Cluster* cluster, GrapplerItem* item) {
     initialize->set_op("Assign");
     initialize->set_device(device);
     (*initialize->mutable_attr())["T"].set_type(dtype);
+    (*initialize->mutable_attr())["use_locking"].set_b(false);
+    (*initialize->mutable_attr())["validate_shape"].set_b(false);
     *initialize->add_input() = tmp_var->name();
     *initialize->add_input() = zeros->name();