worker_devices moving to ds.extended (tensorflow#5775)

hotbbsun · Nov 21, 2018 · b2522f9 · b2522f9
1 parent 9a4848a
commit b2522f9
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/official/utils/misc/distribution_utils_test.py b/official/utils/misc/distribution_utils_test.py
@@ -28,20 +28,20 @@ class GetDistributionStrategyTest(tf.test.TestCase):
   def test_one_device_strategy_cpu(self):
     ds = distribution_utils.get_distribution_strategy(0)
     self.assertEquals(ds.num_replicas_in_sync, 1)
-    self.assertEquals(len(ds.worker_devices), 1)
-    self.assertIn('CPU', ds.worker_devices[0])
+    self.assertEquals(len(ds.extended.worker_devices), 1)
+    self.assertIn('CPU', ds.extended.worker_devices[0])
 
   def test_one_device_strategy_gpu(self):
     ds = distribution_utils.get_distribution_strategy(1)
     self.assertEquals(ds.num_replicas_in_sync, 1)
-    self.assertEquals(len(ds.worker_devices), 1)
-    self.assertIn('GPU', ds.worker_devices[0])
+    self.assertEquals(len(ds.extended.worker_devices), 1)
+    self.assertIn('GPU', ds.extended.worker_devices[0])
 
   def test_mirrored_strategy(self):
     ds = distribution_utils.get_distribution_strategy(5)
     self.assertEquals(ds.num_replicas_in_sync, 5)
-    self.assertEquals(len(ds.worker_devices), 5)
-    for device in ds.worker_devices:
+    self.assertEquals(len(ds.extended.worker_devices), 5)
+    for device in ds.extended.worker_devices:
       self.assertIn('GPU', device)