experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add new kwargs key to escape from creating mirrored variable in TPU strategy scope.

Browse Source 
PiperOrigin-RevId: 272926391
This commit is contained in:
A. Unique TensorFlower 2019-10-04 11:53:06 -07:00 committed by TensorFlower Gardener
parent fe1178ab88
commit 20ebd8a796
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/python/distribute/tpu_strategy.py
View File

@ -386,6 +386,9 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
def _create_variable(self, next_creator, *args, **kwargs):
"""Create a TPUMirroredVariable. See `DistributionStrategy.scope`."""
if kwargs.pop("tpu_embedding_variable_creator", False):
return next_creator(*args, **kwargs)
colocate_with = kwargs.pop("colocate_with", None)
if colocate_with is None:
device_map = self._device_map