diff options
author | 2018-07-31 23:18:58 -0700 | |
---|---|---|
committer | 2018-07-31 23:23:48 -0700 | |
commit | 26ba623dccacfb2f913951e12089a8340e6a11ac (patch) | |
tree | bfa67abb9eb1d408bb1a4a52b8c7c5c3d470b00f /tensorflow/compiler/xla/service/batchnorm_expander_test.cc | |
parent | 92279f8bfa6ce2124439aabfa6db84d722dc2b66 (diff) |
Cleanup the sharding unique device API.
PiperOrigin-RevId: 206885051
Diffstat (limited to 'tensorflow/compiler/xla/service/batchnorm_expander_test.cc')
-rw-r--r-- | tensorflow/compiler/xla/service/batchnorm_expander_test.cc | 6 |
1 files changed, 3 insertions, 3 deletions
diff --git a/tensorflow/compiler/xla/service/batchnorm_expander_test.cc b/tensorflow/compiler/xla/service/batchnorm_expander_test.cc index 32f785a70a..a725351462 100644 --- a/tensorflow/compiler/xla/service/batchnorm_expander_test.cc +++ b/tensorflow/compiler/xla/service/batchnorm_expander_test.cc @@ -137,9 +137,9 @@ ENTRY entry { if (instruction->opcode() == HloOpcode::kParameter) { continue; } - ASSERT_TRUE(instruction->has_sharding()); - TF_ASSERT_OK_AND_ASSIGN(int device, instruction->sharding().UniqueDevice()); - EXPECT_EQ(device, 1); + auto device = instruction->sharding_unique_device(); + ASSERT_TRUE(device); + EXPECT_EQ(*device, 1); } } |