Replace usages of Tensorflow DistributionStrategy method experimental_run_v2 with run.
PiperOrigin-RevId: 318852933 Change-Id: I10e255518dbed9bff095dc08d46d6cb8dde58830
This commit is contained in:
parent
e085901e2c
commit
f96a4f7478
tensorflow/python/distribute
@ -83,7 +83,7 @@ def call_for_each_replica(strategy, fn, args=None, kwargs=None):
|
||||
"overhead currently. We will be working on improving "
|
||||
"this in the future, but for now please wrap "
|
||||
"`call_for_each_replica` or `experimental_run` or "
|
||||
"`experimental_run_v2` inside a tf.function to get "
|
||||
"`run` inside a tf.function to get "
|
||||
"the best performance." % strategy.__class__.__name__, 5)
|
||||
else:
|
||||
# When a tf.function is wrapped to trigger _call_for_each_replica (see
|
||||
|
@ -57,7 +57,7 @@ class AggregatingVariableTest(test.TestCase, parameterized.TestCase):
|
||||
|
||||
per_replica_results = self.evaluate(
|
||||
distribution.experimental_local_results(
|
||||
distribution.experimental_run_v2(assign)))
|
||||
distribution.run(assign)))
|
||||
self.assertAllEqual([3], per_replica_results)
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user