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

Internal change

Browse Source 
PiperOrigin-RevId: 312090528
Change-Id: I474709513b01db8c24c50fd670029451c51cb622
This commit is contained in:
Rohan Jain 2020-05-18 09:08:29 -07:00 committed by TensorFlower Gardener
parent 55aee9e550
commit 46f7108d78
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/python/keras/layers/embeddings.py
View File

@ -129,8 +129,10 @@ class Embedding(Layer):
# since it knows all kernels using the variable only exist on CPU.
# When eager execution is enabled, the placement decision has to be made
# right now. Checking for the presence of GPUs to avoid complicating the
# TPU codepaths which can handle sparse optimizers.
if context.executing_eagerly() and context.context().num_gpus():
# TPU codepaths which can handle sparse optimizers. But if we are within
# a tf.function, we go back the graph mode logic and rely on the placer.
if (context.executing_eagerly() and context.context().num_gpus() and
not ops.inside_function()):
with ops.device('cpu:0'):
self.embeddings = self.add_weight(
shape=(self.input_dim, self.output_dim),