Simplify default key for root tensors.

riga · Dec 23, 2016 · a94863a · a94863a
1 parent 5a5f1fd
commit a94863a
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tfdeploy.py b/tfdeploy.py
@@ -153,9 +153,10 @@ def add(self, tensor, tf_sess=None, key=None, **kwargs):
             tensor = Tensor(tensor, tf_sess, **kwargs)
 
         if key is None:
-            key = len(self.roots)
-            while key in self.roots:
-                key += 1
+            if len(self.roots) == 0:
+                key = 0
+            else:
+                key = max(self.roots.keys()) + 1
 
         self.roots[key] = tensor