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

Merge pull request from geetachavan1/cherrypicks_F4YCK

Browse Source 
[CherryPick 2.3] Going back to forcing embedding layer variables on the CPU even within a tf.function as this is breaking some user code.
This commit is contained in:
Goldie Gadde 2020-07-22 14:12:10 -07:00 committed by GitHub
commit ca2b7ba75c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 2 deletions

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

@ -132,8 +132,7 @@ class Embedding(Layer):
# right now. Checking for the presence of GPUs to avoid complicating the
# 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()):
if context.executing_eagerly() and context.context().num_gpus():
with ops.device('cpu:0'):
self.embeddings = self.add_weight(
shape=(self.input_dim, self.output_dim),