To pass test, add VariableV2 in _PSDeviceChooser

envin3 · Jun 5, 2017 · 01db54e · 01db54e
1 parent 2fa4bb6
commit 01db54e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/slim/deployment/model_deploy.py b/slim/deployment/model_deploy.py
@@ -663,7 +663,7 @@ def choose(self, op):
         if op.device:
           return op.device
         node_def = op if isinstance(op, tf.NodeDef) else op.node_def
-        if node_def.op == 'Variable':
+        if node_def.op.startswith('Variable'):
           t = self._task
           self._task = (self._task + 1) % self._tasks
           d = '%s/task:%d' % (self._device, t)