diff options
Diffstat (limited to 'tensorflow/core/grappler/optimizers/model_pruner_test.cc')
| -rw-r--r-- | tensorflow/core/grappler/optimizers/model_pruner_test.cc | 10 |
1 files changed, 5 insertions, 5 deletions
diff --git a/tensorflow/core/grappler/optimizers/model_pruner_test.cc b/tensorflow/core/grappler/optimizers/model_pruner_test.cc index aea1fcd7c9..ee722f311e 100644 --- a/tensorflow/core/grappler/optimizers/model_pruner_test.cc +++ b/tensorflow/core/grappler/optimizers/model_pruner_test.cc @@ -320,14 +320,14 @@ TEST_F(ModelPrunerTest, PruningPerservesCrossDeviceIdentity) { Output c = ops::Const(s.WithOpName("c").WithDevice("/cpu:0"), 0.0f, {10, 10}); // Node i1 should be preserved. - Output i1 = ops::Identity(s.WithOpName("i1").WithDevice("/gpu:0"), c); - Output a1 = ops::Sqrt(s.WithOpName("a1").WithDevice("/gpu:0"), {i1}); - Output a2 = ops::Sqrt(s.WithOpName("a2").WithDevice("/gpu:0"), {i1}); + Output i1 = ops::Identity(s.WithOpName("i1").WithDevice("/device:GPU:0"), c); + Output a1 = ops::Sqrt(s.WithOpName("a1").WithDevice("/device:GPU:0"), {i1}); + Output a2 = ops::Sqrt(s.WithOpName("a2").WithDevice("/device:GPU:0"), {i1}); // Node i2 should be pruned since it resides on the sender's device. Output i2 = ops::Identity(s.WithOpName("i2").WithDevice("/cpu:0"), c); - Output a3 = ops::Sqrt(s.WithOpName("a3").WithDevice("/gpu:0"), {i2}); - Output a4 = ops::Sqrt(s.WithOpName("a4").WithDevice("/gpu:0"), {i2}); + Output a3 = ops::Sqrt(s.WithOpName("a3").WithDevice("/device:GPU:0"), {i2}); + Output a4 = ops::Sqrt(s.WithOpName("a4").WithDevice("/device:GPU:0"), {i2}); GrapplerItem item; TF_CHECK_OK(s.ToGraphDef(&item.graph)); |