Improve performance of MapGradient

tensorpack/tfutils/gradproc.py

@@ -128,7 +128,8 @@ class MapGradient(GradientProcessor):
         for grad, var in grads:
             if re.match(self.regex, var.op.name):
                 matched = True
-                grad = self.func(grad, var)
+                with tf.device(grad.device):
+                    grad = self.func(grad, var)
                 if grad is not None:
                     ret.append((grad, var))
             else: