From d81bd790cb18a23795a6d0b83d683071c77ca6d2 Mon Sep 17 00:00:00 2001 From: Vijay Vasudevan Date: Fri, 4 Mar 2016 22:03:52 -0800 Subject: TensorFlow: Change apply_dense and apply_sparse to use a colocation constraint rather than ops.device, since colocation is more portable. Change: 116431514 --- tensorflow/python/training/optimizer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/training/optimizer.py b/tensorflow/python/training/optimizer.py index 9af92c66bf..1e8d6b0f12 100644 --- a/tensorflow/python/training/optimizer.py +++ b/tensorflow/python/training/optimizer.py @@ -292,7 +292,9 @@ class Optimizer(object): for grad, var in grads_and_vars: if not grad: continue - with ops.name_scope("update_" + var.op.name), ops.device(var.device): + # We colocate all ops created in _apply_dense or _apply_sparse + # on the same device as the variable. + with ops.name_scope("update_" + var.op.name), ops.colocate_with(var): if isinstance(grad, ops.Tensor): update_ops.append(self._apply_dense(grad, var)) else: -- cgit v1.2.3