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

Add additional constraint to place identity under CPU

Browse Source 
Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
This commit is contained in:
Yong Tang 2020-05-22 19:12:20 +00:00
parent 4aecaf3771
commit ee92f24165
1 changed files with 2 additions and 1 deletions

3
tensorflow/python/ops/array_grad.py
View File

@ -641,7 +641,8 @@ def _GatherV2Grad(op, grad):
# For axis 0 gathers, build an appropriately shaped IndexedSlices.
if axis_static == 0:
if context.executing_eagerly():
params_tail_shape = array_ops.identity(params_shape)[1:]
with ops.device("/cpu:0"):
params_tail_shape = array_ops.identity(params_shape)[1:]
else:
params_tail_shape = params_shape[1:]
values_shape = array_ops.concat([indices_size, params_tail_shape], 0)